python遍历n叉树

Question

传入参数为一个二层嵌套list(例：[[1,2],['a','b','c'],['!','@'],[...]])，list[0]为第一层节点，list[1]为第二层节点，依此类推。

要求：设计一个函数，return该n叉树所有的路径。
例子中的list就是return[
[1,'a','!'],
[1,'a','@'],
[1,'b','!'],
[1,'b','@'],
[1,'c','!'],
[1,'c','@'],
[2,'a','!'],
[2,'a','@'],
[2,'b','!'],
[2,'b','@'],
[2,'c','!'],
[2,'c','@']]
如下图所示：

Answer 1

这个并不是树的遍历，这就是一个简单穷举的问题，第一个位置可能是1或2，第二个可能是a，b或c，以此类推。

def enumerateAll(levels):
    if len(levels) is 1:
        yield from map(lambda e: [ e ], levels[0])
    else:
        head = levels[0]
        tail = levels[1:]
        for h in head:
            for t in enumerateAll(tail):
                yield [ h ] + t

list(enumerateAll([ [ 1, 2 ], [ 'a', 'b', 'c' ], [ '!', '@' ] ]))

应该有更简洁的写法，暂时不想想了

Answer 2

一时短路差点忘了笛卡尔乘积，最简单，两行代码。

from itertools import product
[list(i) for i in product(*tree)]

-------------------update ----------------------
我有一个使用reduce的写法，只是不知道是否好理解

from functools import reduce

tree = [[1,2],['a','b','c'],['!','@']]

tmp_tree = [[[]]] + tree

def step(x, y):
    return [i + [j]  for i in x for j in y]

reduce(step, tmp_tree)

tmp_tree 是为了方便reduce同构迭代，当然你也可以这样写

from functools import reduce

tree = [[1,2],['a','b','c'],['!','@']]

def step(x, y):
    return [i + [j] if isinstance(i, list) else [i, j] for i in x for j in y]

reduce(step, tree)