dheixis
dheixis

Reputation: 23

How to enumerate the paths to every element in an arbitrarily deep n-ary tree

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

Answers (2)

Lante Dellarovere
Lante Dellarovere

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

Patrick Haugh
Patrick Haugh

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

Related Questions