Reputation: 23
I'm working on some Python code where I have a representation of data as n-ary trees where n is given by the user. The arity of the tree is definite but I'm struggling with an algorithm to enumerate the path from the root to each element.
For example, if I have a 3-ary tree
.
/|\
/ | \
/ | \
/ | \
. . .
/|\ /|\ /|\
a b cd . hi j k
/|\
e f g
represented by the nested list
[[a, b, c], [d, [e, f, g], h], [i, j, k]]
I'd like to get a list of tuples like
[(a, 00), (b, 01), (c, 02), (d, 10), (e, 110), (f, 111), (g, 112), (h, 12), (i, 20), (j, 21), (k, 22)]
I did find a similar problem here Enumerating all paths in a tree but it's not quite what I need and I'm not sure how I might achieve the kind of enumeration I'm looking for.
Upvotes: 0
Views: 180
Reputation: 1858
I think there is a mismatch between actual tree and its representation.
If I didn't mess up with the picture it should be:
repr = [["a", "b", "c"], ["d", ["e", "f", "g"], "h"], ["i", "j", "k"]]
If your data are composed by lists like repr
, can use a recursive function like this:
def tree(l, ind=""):
for i, x in enumerate(l):
if isinstance(x, list):
yield from tree(x, ind + str(i))
else:
yield x, ind + str(i)
>>> print(list(tree(repr))
[('a', '00'), ('b', '01'), ('c', '02'), ('d', '10'), ('e', '110'), ('f', '111'), ('g', '112'), ('h', '12'), ('i', '20'), ('j', '21'), ('k', '22')]
Upvotes: 1
Reputation: 60994
Here's a recursive method to bubble up the path to each leaf using yield from
and recursion.
class Tree:
def __init__(self, *children, data=None):
self.data = data
self.children = children
def find_all(root, path_to=()):
if root is None:
return
if not root.children:
yield (root.data, path_to)
else:
for i, node in enumerate(root.children):
yield from find_all(node, path_to=(*path_to, i))
root = Tree(Tree(Tree(data='a'), Tree(data='b'), Tree(data='c')), Tree(Tree(data='d'), Tree(Tree(data='e'), Tree(data='f'), Tree(data='g')), Tree(data='f')), Tree(Tree(data='g'), Tree(data='h'), Tree(data='i')))
print(list(find_all(root)))
# [('a', (0, 0)), ('b', (0, 1)), ('c', (0, 2)), ('d', (1, 0)), ('e', (1, 1, 0)), ('f', (1, 1, 1)), ('g', (1, 1, 2)), ('f', (1, 2)), ('g', (2, 0)), ('h', (2, 1)), ('i', (2, 2))]
Upvotes: 0