Calculate height of an arbitrary (non-binary) tree

Question

I'm currently taking on online data structures course and this is one of the homework assignments; please guide me towards the answer rather than giving me the answer.

The prompt is as follows:

Task. You are given a description of a rooted tree. Your task is to compute and output its height. Recall that the height of a (rooted) tree is the maximum depth of a node, or the maximum distance from a leaf to the root. You are given an arbitrary tree, not necessarily a binary tree.

Input Format. The first line contains the number of nodes n. The second line contains integer numbers from −1 to n−1 parents of nodes. If the i-th one of them (0 ≤ i ≤ n−1) is −1, node i is the root, otherwise it’s 0-based index of the parent of i-th node. It is guaranteed that there is exactly one root. It is guaranteed that the input represents a tree.

Constraints. 1 ≤ n ≤ 10⁵.

My current solution works, but is very slow when n > 10². Here is my code:

# python3

import sys
import threading

# In Python, the default limit on recursion depth is rather low,
# so raise it here for this problem. Note that to take advantage
# of bigger stack, we have to launch the computation in a new thread.
sys.setrecursionlimit(10**7)  # max depth of recursion
threading.stack_size(2**27)   # new thread will get stack of such size
threading.Thread(target=main).start()

# returns all indices of item in seq
def listOfDupes(seq, item):
    start = -1
    locs = []
    while True:
        try:
            loc = seq.index(item, start+1)
        except:
            break
        else:
            locs.append(loc)
            start = loc
    return locs

def compute_height(node, parents):
    if node not in parents:
        return 1
    else:
        return 1 + max(compute_height(i, parents) for i in listOfDupes(parents, node))

def main():
    n = int(input())
    parents = list(map(int, input().split()))
    print(compute_height(parents.index(-1), parents))

Example input:
>>> 5
>>> 4 -1 4 1 1
This will yield a solution of 3, because the root is 1, 3 and 4 branch off of 1, then 0 and 2 branch off of 4 which gives this tree a height of 3.

How can I improve this code to get it under the time benchmark of 3 seconds? Also, would this have been easier in another language?

Answer 1

Python will be fine as long as you get the algorithm right. Since you're only looking for guidance, consider:

1) We know the depth of a node iif the depth of its parent is known; and 2) We're not interested in the tree's structure, so we can throw irrelevant information away.

The root node pointer has the value -1. Suppose that we replaced its children's pointers to the root node with the value -2, their children's pointers with -3, and so forth. The greatest absolute value of these is the height of the tree.

If we traverse the tree from an arbitrary node N(0) we can stop as soon as we encounter a negative value at node N(k), at which point we can replace each node with the value of its parent, less one. I.e, N(k-1) = N(k) -1, N(k-2)=N(k-1) - 1... N(0) = N(1) -1. As more and more pointers are replaced by their depth, each traversal is more likely to terminate by encountering a node whose depth is already known. In fact, this algorithm takes basically linear time.

So: load your data into an array, start with the first element and traverse the pointers until you encounter a negative value. Build another array of the nodes traversed as you go. When you encounter a negative value, use the second array to replace the original values in the first array with their depth. Do the same with the second element and so forth. Keep track of the greatest depth you've encountered: that's your answer.

Answer 2

The structure of this question looks like it would be better solved bottom up rather than top down. Your top-down approach spends time seeking, which is unnecessary, e.g.:

def height(tree):
    for n in tree:
        l = 1
        while n != -1:
            l += 1
            n = tree[n]
        yield l

In []:
tree = '4 -1 4 1 1'
max(height(list(map(int, tree.split()))))

Out[]:
3

Or if you don't like a generator:

def height(tree):
    d = [1]*len(tree)
    for i, n in enumerate(tree):
        while n != -1:
            d[i] += 1
            n = tree[n]
    return max(d)

In []:
tree = '4 -1 4 1 1'
height(list(map(int, tree.split())))

Out[]:
3

The above is brute force as it doesn't take advantage of reusing parts of the tree you've already visited, it shouldn't be too hard to add that.

Answer 3

Your algorithm spends a lot of time searching the input for the locations of numbers. If you just iterate over the input once, you can record the locations of each number as you come across them, so you don't have to keep searching over and over later. Consider what data structure would be effective for recording this information.