Searching any tree in Python

I need to write a function in Python that takes a tree and an index and returns the subtree or leaf at that index.

I tried with loops and nested loops until I realized that the tree that had to run the tests was always the same:

tree = (((1, 2), 3), (4, (5, 6)), 7, (8, 9, 10))

which actually looks like this:

Sample tree

So all I needed to pass the tests given was this:

def tree_ref(tree, index):
    if len(index) == 1:
        print tree[index[0]]
    elif len(index) == 2:
        print tree[index[0]][index[1]]
    elif len(index) == 3:
        print tree[index[0]][index[1]][index[2]]
    else:
        return False

For example, for an index = (1, 1, 0) it should return the number 5 and it does.

However, I know this code won’t work for other trees or index with more than 3 elements. How could I make it work for any given tree and index?

Answer

You should try to use recursion.

Something like below:

def tree_ref(tree, index):
    if len(index) == 1:
        print tree[index[0]]
    else:
        tree_ref(tree[index[0]], index[1:])