在Python中返回树的分支作为列表

时间:2017-01-22 15:25:46

标签: python function class tree

我正在尝试在Python中编写一个递归函数,在给定分支的深度或max_sum的情况下,将树的分支作为列表返回。我真的很沮丧。也许有更简单的类或生成器实现?下面是我想要实现的功能行为的详细描述。

func(data, depth)
'''Accepts a list with numbers > 0 and depth, i.e. max elements per list; 
   returns each branch of a tree'''    

----------Examples--------------
Input:  func([2, 1], depth=2)
Output: [[2, 2], [2, 1], [1, 2], [1, 1]]

Input:  func([3, 2, 1], depth=2)
Output: [[3, 3], [3, 2], [3, 1]
         [2, 3], [2, 2], [2, 1]
         [1, 3], [1, 2], [1, 1]]

Input:  func([2, 1], depth=3)
Output: [[2, 2, 2], [2, 2, 1], [2, 1, 2], [2, 1, 1],
         [1, 2, 2], [1, 2, 1], [1, 1, 2], [1, 1, 1]]

第二个例子的图片

第三个例子的图片

这是我写的代码,它只适用于第一个例子,它太可怕了,我真的为此感到羞耻:/我尝试了几十种使用类和生成器的方法,但我对这些和代码并不是很熟悉即使是第一个例子,也只返回了一半的选项。

tree = []
node_list = [2, 1]

def make_branch(depth=2, branch=None, d={0:2, 1:1}, switch=False, count=0):
    #print(count)

    if branch is None:
        branch = []

    for i in range(2):
        #print(i)
        if switch:
            branch.append(d[i+1]) 
            switch=False
        else:
            branch.append(d[i])

        if len(branch) >= depth:
            tree.append(branch)
            print(branch)
            return

        make_branch(count= count + 1, branch=branch)
        #print(-count)
        branch = branch[:-1]


for i in range(len(node_list)):
    if i % 2 == 0:
        make_branch()
    else:
        make_branch(switch=True)

print(tree)

1 个答案:

答案 0 :(得分:0)

我不明白你为什么想把它与穿越一棵树联系起来。您的任务基本上只是生成所有排列(替换) - 这与具有固定集合的笛卡尔积相同 - 在一组数字上给定长度。

在Python中,您可以按照以下方式执行此操作:

import itertools
for i in itertools.product([1,2], repeat=3):
  print i

这将是例如输出你的第三个例子。请注意,每个输出都是一个元组而不是列表 - 所以你可能想要转换它们。

最简单的实现可能会这样:

def prod(lst, depth, buf=''):
    if depth == 0:
        print buf
        return
    for elem in lst:
        prod(lst, depth - 1, buf + str(elem))

prod([1,2], 3)
print 
prod([1,2,3], 2)

输出:

111
112
121
122
211
212
221
222

11
12
13
21
22
23
31
32
33