python - 在 Python 中高效迭代任意深度字典树

Question

我在字典中存储了以下树数据结构:

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 from 和partial 正在减慢我的速度。所以我试着把它弄平:

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 上，主要的减速是回调的结果报告:

• yielding 结果是最慢的 - 与内联相比损失约 30%。循环版本见iterate_tree6，递归版本见iterate_tree3
• 通过从回调中调用回调进行报告稍微好一些——比内联(在 Python3.4 上)慢 17%。请参阅 iterate_tree3_noyield
• 没有任何报告能比内联报告运行得更好。请参见iterate_tree6_nofeedback

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

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

Answer 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:]