Problem
Description: You are given a description of an arbitrary, rooted tree. Your task is to compute and output its height.
Input: The first line contains the number of nodes 𝑛. The second line contains 𝑛the parents of nodes: integer numbers from −1 to 𝑛 − 1. If the 𝑖-th one of them (0 ≤ 𝑖 ≤ 𝑛 − 1) is −1, node 𝑖 is the root. Otherwise, it’s the zero-based index of the parent of the 𝑖-th node. It is guaranteed that there is exactly one root. It is guaranteed that the input represents a tree.
Constraints: 1 ≤ 𝑛 ≤ 10⁵
Output: The height of the tree.
Sample:
INPUT5
4 -1 4 1 1OUTPUT3
The input 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. The height of this tree is 3 because the number of vertices on the path from root 1 to leaf 2 is 3.
Solution
We can break up our task into two main components: building the tree, and computing its height.
Setting Up
We can take advantage of the fact that the nodes of the tree are labeled by contiguous indices from 0 to 𝑛 − 1, and format our tree by just using a two-dimensional array with n rows. Each row would represent a tree node, and the columns would represent the node’s children.
Using the example provided above, we can construct this tree as:
[[], [3, 4], [], [], [0, 2]]Node 1 has children nodes of 3 and 4. Notice how the nested arrays can contain any number of children, so we can format any arbitrary tree as opposed to only binary trees.
Keeping this in mind, we can define a class of TreeHeight , with a function read(self) that reads our input and constructs our class variables.
class TreeHeight:
def read(self):
# n = number of tree nodes
self.n = int(sys.stdin.readline())# parent = list of input
self.parent = list(map(int, sys.stdin.readline().split()))# nodes = 2d array of tree nodes
self.nodes = None# root = the root of the tree
self.root = None
Building the Tree
In TreeHeight, let’s define a function called build_tree that will properly initialize our self.nodes object.
def build_tree(self):# initialize nodes as a 2-dimensional array of n rows
self.nodes = [ [] for i in range(self.n) ]# for each node from 0 to n-1
for child_index in range(self.n):
# get their parent
parent_index = self.parent[child_index]
# if the parent is -1, then the child_index must be the root
if parent_index == -1:
self.root = child_index # otherwise, add the child's index to the parent node
else:
self.nodes[parent_index].append(child_index)
Computing Tree Height
Now that our tree has been constructed, let’s define a recursive function to get the height of the tree. To help with comprehension, think of recursion as defining one or two corner cases, and then leaving the rest of the work to the computer.
# current = the node we are currently on
# height = the height of the tree at the current node
def get_max_height(self, current, height): # if the current node has no children
if not self.nodes[current]: # return the height as it is
return height # otherwise, find the max of each child's heights
max_height = 0 # for each child of the current node
for c in self.nodes[current]: # compare the current max_height with the child's height
max_height = max(max_height, self.get_max_height(c, height+1)) # return max height
return max_height
Execution
Now that we have each component constructed, let’s tie them all together. Define the functioncompute-height that constructs the tree and finds its height.
def compute_height(self):
# build the tree
self.build_tree() # find height by using the root index & an initial height of 1
return self.get_max_height(self.root, 1)
Finally, outside of our class TreeHeight , we can define a main function that reads in our input and computes the tree’s height.
def main():
tree = TreeHeight()
tree.read()
print(tree.compute_height())Conclusion
Here is the final program in its entirety: