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

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))

This runs 3.5 seconds average using Python 3.4.2 on my laptop.
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))

这很通用,但速度慢了两倍! 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))

这是泛型方法,它可以工作,但与递归相比性能有所提升,即仍然比内联版本慢两倍。
那么如何以泛型方式实现我的遍历?内联版本为什么会快那么多?
附言:上面的代码适用于Python 3.3+。我已经改编成Python 2,结果类似。
解决方案和分析
我对所有解决方案和优化进行了比较分析。代码和结果可从 the gist获得。
简而言之,最快的解决方案是使用经过优化的基于循环的版本:
- 它是最快的支持从回调中方便地报告结果的版本 - 它只比内联版本慢30%(在Python3.4上) - 在PyPy上运行时,其速度得到了巨大的提升,甚至超过了内联版本
基于循环的迭代在PyPy上拥有一切。
在非pypy上,主要减速是由于从回调中报告结果:
  • yield 的结果最慢,比内联慢约30%。请参见 iterate_tree6(循环版本)和iterate_tree3(递归版本)
  • 通过从回调中调用回调进行报告略好于内联 - 在Python3.4上比内联慢17%。请参见iterate_tree3_noyield
  • 完全不报告可能比内联更好。请参见iterate_tree6_nofeedback

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

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


结果对于Python 2来说不应该是相似的,至少对于大型字典而言如此--在Python 2中,items()返回一份键值对的副本,而在Python 3中它返回一个“视图”,而不复制任何内容。 - jedwards
你说得对。但是我的树是“正方形”的,所以它不会受到太大的影响。 - Zaar Hai
3个回答

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都会向回调中添加一个额外的函数调用层次。


你的版本确实比我这个“iterate_tree”快大约25%,这是非常好的提升。谢谢。Gist链接:https://gist.github.com/haizaar/a3cda2dc17a814f08aa5 - Zaar Hai
哈哈,你的第一条评论让我有一段时间非常困惑! - Andrew Magee

2

这是一个经过优化的迭代版本的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:]

1
太棒了!实际上与安德鲁的代码一样。另一个好处是在PyPy上,基于循环的解决方案获得了惊人的提升。因此,我认为您的解决方案是最通用的。测试Gist:https://gist.github.com/haizaar/a3cda2dc17a814f08aa5(您的是iterate_tree6) - Zaar Hai

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, [])

Edit操作采用类似的方法,但根据您的评论保留了关键字

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

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

它正在累加回调函数的值,而我需要累加回调函数的键。所以你的函数做错了事情。并且它没有提供一种消耗回调函数结果的方式(就像我用yield那样)。 - Zaar Hai
@ZaarHai,除非您在回调中需要使用路径键,否则将字典的引用作为回调参数会更有效率,而不是根据您累积的键重新遍历。这是情况吗?您需要这些键吗? - jedwards
是的,我绝对需要这些键。它们是“路径”。回调函数需要接收 (key1, key2, ..., keyN, val_of_last_dict) 并且每个叶子节点都会被调用一次。 - Zaar Hai
@ZaarHai更新了一个类似的方法,它可以跟踪键。看看你的基准测试结果在哪里会很有趣。 - jedwards
这是结果:https://gist.github.com/haizaar/a3cda2dc17a814f08aa5。看起来,从回调函数返回结果是导致速度变慢的原因。最慢的方法是 yield,从回调函数中调用回调函数更好。完全不报告结果提供了类似于使用 @Andew 版本的内联结果(您的稍微慢一些)。无论如何,谢谢! - Zaar Hai

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接