在Python中有效地迭代任意深度字典树

时间:2015-03-19 08:38:09

标签: python algorithm tree python-3.4

我将以下树数据结构存储在词典中:

1
   2
      3
         4 -> ["a", "b", "c"]
         5 -> ["x", "y", "z"]
   3
      5
         7 -> ["e", "f", "j"]

以下是我在Python中构建它的示例:

tree = dict()
for i in range(100):
    tree[i] = dict()
    for j in range(10):
        tree[i][j] = dict()
        for k in range(10):
            tree[i][j][k] = dict()
            for l in range(10):
                tree[i][j][k][l] = dict()
                for m in range(10):
                    tree[i][j][k][l][m] = dict()
                    for n in range(10):
                        tree[i][j][k][l][m][n] = ["a", "b", "c", "d", "e", "f", "g"]

我想遍历它并在到达每个叶子时进行一些计算。在进行计算时,我需要知道叶子的路径。

即。给予回调

def callback(p1, p2, p3, p4, leaf):
    ...

我希望像使用我的树例子一样调用它:

callback(1, 2, 3, 4, ["a", "b", "c"])
callback(1, 2, 3, 5, ["x", "y", "z"])
callback(1, 3, 5, 7, ["e", "f", "j"])

问题:如何最有效地实施遍历?请注意,树深度不是静态的。

以下是我的尝试:

1。内联代码。这是最快的代码,但在实践中无法使用,因为树深度不是静态的。

def callback(*args):
    assert isinstance(args[-1], list)

start = time.time()
for k1, leafs1 in tree.items():
    for k2, leafs2 in leafs1.items():
        for k3, leafs3 in leafs2.items():
            for k4, leafs4 in leafs3.items():
                for k5, leafs5 in leafs4.items():
                    for k6, val in leafs5.items():
                        callback(k1, k2, k3, k4, k5, k6, val)
print("inline: %f" % (time.time() - start))

在笔记本电脑上使用Python 3.4.2平均运行3.5秒。

2。递归方法

from functools import partial
def iterate_tree(tree, depth, callback):
    if depth:
        for k, subtree in tree.items():
            cb = partial(callback, k)
            yield from iterate_tree(subtree, depth-1, cb)
    else:
        for k, v in tree.items():
            rv = callback(k, v)
            yield rv

start = time.time()
for i in iterate_tree(tree, 5, callback):
    pass
print("iterate_tree: %f" % (time.time() - start))

这是通用的,一切都很好,但慢了2倍!

第3。非递归方法我认为这可能是递归,yield frompartial正在减慢我的速度。所以我试着把它打开:

def iterate_tree2(tree, depth, callback):
    iterators = [iter(tree.items())]
    args = []
    while iterators:
        try:
            k, v = next(iterators[-1])
        except StopIteration:
            depth += 1
            iterators.pop()
            if args:
                args.pop()
            continue

        if depth:
            args.append(k)
            iterators.append(iter(v.items()))
            depth -= 1
        else:
            yield callback(*(args + [k, v]))

start = time.time()
for i in iterate_tree2(tree, 5, callback):
    pass
print("iterate_tree2: %f" % (time.time() - start))

这是泛型并且有效,但与递归相比性能提升,即仍然比内联版本慢两倍。

那么如何以通用方式实现遍历?是什么让内联版本更快?

P.S。上面的代码适用于Python 3.3+。我已将它改编为Python 2,结果类似。

解决方案和分析

我对所有解决方案和优化进行了比较分析。代码和结果可以从the gist获得。

TL; DR;最快的解决方案是使用优化的基于循环的版本:

  • 它是支持回调方便结果报告的最快版本
  • 它比内联版本(在Python3.4上)慢了30%
  • 在PyPy上,它获得了惊人的速度提升,甚至超越了内联版本

基于循环的迭代在PyPy上运行时拥有所有内容。

在非小游戏的情况下,主要的减速是回调报告的结果:

    与内联相比,
  • yield结果是最慢的 - 约30%的惩罚。请参阅iterate_tree6获取循环版本,iterate_tree3获取递归版本
  • 通过从回调调用回调进行报告稍微好一些 - 比内联慢了17%(在Python3.4上)。请参阅iterate_tree3_noyield
  • 根本没有报告可以比内联更好地运行。请参阅iterate_tree6_nofeedback

对于基于递归的版本,使用元组进行参数累积而不是列表。性能差异非常显着。

感谢所有为此话题做出贡献的人。

3 个答案:

答案 0 :(得分:2)

我设法将内联版本和第一个递归版本之间的性能提高到了一半,我认为是等效的。

def iterate_tree_2(tree, depth, accumulator, callback):
    if depth:
        for k, subtree in tree.items():
            yield from iterate_tree_2(subtree, depth-1, accumulator + (k,), callback)
    else:
        for k, v in tree.items():
            rv = callback(accumulator + (k,), v)
            yield rv

>>> for i in iterate_tree_2(tree, depth, (), callback): pass

它与

调用回调略有不同
callback((1, 2, 3, 4), ["a", "b", "c"])

而不是

callback(1, 2, 3, 4, ["a", "b", "c"])

实现的不同之处在于它构建了参数元组,而不是使用partial。我猜这是有意义的,因为每次你致电partial时,你都会在回调中添加额外的函数调用层。

答案 1 :(得分:1)

这是一种递归方法,似乎比你的内联方法执行大约5-10%更好

def iter_tree(node, depth, path):
    path.append(node)
    for v in node.values():
        if depth:
            iter_tree(v, depth-1, path)
        else:
            callback(path)

您可以致电:

iter_tree(tree, 5, [])
根据您的评论,

修改类似方法,但保留了

def iter_tree4(node, depth, path):
    for (k,v) in node.items():
        kpath = path + [k]
        if depth:
            iter_tree4(v, depth-1, kpath)
        else:
            callback(kpath, v)

以同样的方式打电话。

请注意,我们仅通过跟踪值而失去了性能提升,但它仍然与您的内联方法竞争:

Iteration 1  21.3142
Iteration 2  11.2947
Iteration 3   1.3979

列出的数字是性能损失百分比:[(递归内联)/内联]

答案 2 :(得分:1)

这是迭代iterate_tree2的优化版本。它在我的系统上快了40%,主要得益于改进的循环结构和try except的消除。 Andrew Magee的递归代码大致相同。

def iterate_tree4(tree, depth, callback):
    iterators = [iter(tree.items())]
    args = [] 
    while iterators:
        while depth:
            for k, v in iterators[-1]:
                args.append(k)
                iterators.append(iter(v.items()))
                depth -= 1
                break
            else:
                break
        else:
            for k, v in iterators[-1]:
                yield callback(*(args + [k, v]))
        depth += 1
        del iterators[-1]
        del args[-1:]