CODEWITHSUNDEEP
pythonalgorithmtreecombinations
Ausrada404
Ausrada404

Reputation: 599

How to traverse the tree to get all combinations of the leaf node?

I have a tree like the below and I want to get all combinations of the leaf node in the penultimate layer:

enter image description here

[E,C,H]
[E,C,I]
[E,C,J]
[E,C,K]
[F,C,H]
[F,C,I]
[F,C,J]
[F,C,K]

I want to use a function like product to implement this traverse.

The code to create the tree:

class Node:
    def __init__(self, name, parent=None):
        self.name = name
        self.parent = parent
        self.children = []

    def add(self, child):
        self.children.append(child)


    def showDetails(self, depth=1):
        print(self.name)
        for child in self.children:
            print("\t"*depth, end="")
            child.showDetails(depth+1)

a = Node('A')
b = Node('B', a)
c = Node('C', a)
d = Node('D', a)
a.add(b)
a.add(c)
a.add(d)

e = Node('E', b)
f = Node('F', b)
b.add(e)
b.add(f)

g = Node('G', d)
d.add(g)

h = Node('H', g)
i = Node('I', g)
j = Node('J', g)
k = Node('K', g)
g.add(h)
g.add(i)
g.add(j)
g.add(k)

a.showDetails()

Upvotes: 1

Views: 293

Answers (1)

Ajax1234
Ajax1234

Reputation: 71461

You can use recursion to build the leaf levels, and then apply itertools.product to produce the combinations:

import itertools
def leaves(tree):
  for x, y in itertools.groupby(tree.children, key=lambda x:not x.children):
     if x:
        yield [i.name for i in y]
     else:
        for i in y:
           yield from leaves(i)

combos = [*itertools.product(*leaves(a))]

Output:

[('E', 'C', 'H'), ('E', 'C', 'I'), ('E', 'C', 'J'), ('E', 'C', 'K'), ('F', 'C', 'H'), ('F', 'C', 'I'), ('F', 'C', 'J'), ('F', 'C', 'K')]

Upvotes: 1

Related Questions