calculate tree height (arbitrary tree) - python

Question

Given an array as input, we have to calculate the tree height. For example, input [4, -1, 4, 1, 1] means that there are 5 nodes with numbers from 0 to 4, node 0 is a child of node 4, node 1 is the root, node 2 is a child of node 4, node 3 is a child of node 1 and node 4 is a child of node 1.

I have a very low performing code that loops over the input more than once:

def height(array):
 di= {array.index(-1):-1}
 while (len(di)<len(array)):
 for i, p in enumerate(array):
 if p in di.keys():
 di[i] = di[p] - 1 
 max_distance = -1
 for val in di.values():
 if val < max_distance:
 max_distance = val
 max_distance = abs(max_distance)
 return max_distance

Another method I'm trying is I group the child nodes together like this: input [9, 7, 5, 5, 2, 9, 9, 9, 2, -1] I group then to {9: [0, 5, 6, 7], 7: [1], 5: [2, 3], 2: [4, 8], -1: [9]} But I'm stuck and have no idea if I'm on the right track or what to do after this. Please let me know what you think. Much appreciated!

• superb rain: Is that testable online somewhere?
• superb rain: Are there limits?
• pacmaninbw: I suggest you draw the trees as graphs to redesign your code.
• Chris: This is from coursera data structures course and the grader output was time limit exceeded. I changed the code according to the advice from the answer and it works!

Answer 1

• Standard indentation is four spaces.
• Missing spaces here and there.
• Unnecessary parentheses after while.
• Can't handle empty trees.
• Name di is really unclear, i and p could also be better.
• Using .keys() like that is pointless.
• Using negative values is pointless and confusing. Using abs on the result is particularly confusing, as it suggests that the value you give it could be negative or positive. Since you know it's negative, that should've been just -max_distance.
• No need to reinvent max.

Rewritten accordingly:

def height(array):
 depth = {-1: 0}
 while len(depth) <= len(array):
 for child, parent in enumerate(array):
 if parent in depth:
 depth[child] = depth[parent] + 1
 return max(depth.values())

For performance, instead of repeatedly searching for children with finished parents, first do one pass to collect each parent's children. Then when you finish a parent, simply go through its previously collected children to finish them, too. This is then linear instead of quadratic time.

Or if the limits are small enough (so far you still haven't answered the comment asking about that), you could use a simple recursive depth(node) function decorated with functools.cache, apply it to all nodes, and take the max. Also linear time.