PROBLEM LINK:
Author: Kunal Khanna
Editorialist: Arjun P
DIFFICULTY:
Easy-Medium
PREREQUISITES:
None
PROBLEM:
You have a complete binary tree of N nodes. Given an integer X, find two quantities:
a)The number of nodes which are at X steps from the root, i.e. the number of nodes such that, starting at that node, you have to take X steps upwards to reach the root.
b)The number of nodes such that the maximum distance from the node to any leaf in the subtree of the node is X.
QUICK EXPLANATION
The first quantity is just the number of nodes at depth X, where depth is 0-indexed. Take care of the cases where the X^{th} depth does not exist and when X refers to the bottom-most row of the tree, which may be partially filled. In this case the answer is N - (2^{h-1} - 1).
Now, for finding the second quantity, note that at a depth d, nodes which have some nodes from the last row in their subtree have maximum distance (h-1) - d), and nodes which don’t, have distance as (h-2) - d). Start from the right-most node at the last row. Keep moving up its parents. Use the fact that nodes which are the current node, or which are to the left of the current node, always contain the last row in their subtree, and nodes to the right do not contain it.
EXPLANATION
Note that the first quantity is just the number of nodes at a depth X from the root.
Observe that, in a complete binary tree, depth 0 has 1 node, each of these has 2 children, and so, depth 1 has 2 nodes, and depth 2 has 4 nodes, and so on. each depth has twice the number of nodes of the previous depth because each node at that depth has two children each. Thus depth i has 2^i nodes. This is true for all depths except the last; the last depth may be partially filled.
First, find h, the height of the complete binary tree. The height of the tree is given by
h = 1 + \lfloor \log_2 N \rfloor
Note that depth is 0-indexed and depth (h-1) is the last depth.
If X > (h - 1), then no solution exists and answer is zero.
If X < (h - 1), then answer is 2^X, since the row is fully filled. (every row except the last row is fully filled)
If X = (h - 1), then we need to find the number of nodes on the last row. This is just N minus the number of elements in all the other rows, i.e.
ans = N - 2^0 - 2^1 \ldots -2^{h-2} = N - (2^{h-1} - 1)
Now on to the second part. We have to find the number of nodes such that the maximum distance from the node to any leaf in the subtree of the node is X. If the last row was fully filled out, we would again just have to find the X^{th} depth from the bottom. However, since the last row may not be fully filled, we cannot do this. This is because, on any given row, some nodes will have some nodes from the last row in their subtrees, and some will not. Clearly, the nodes that do not have the last row in their subtrees will always have the second last row in their subtree, and the maximum distance is the distance to the leaf nodes on the second last row. For nodes which have nodes from the last row in their subtree, the maximum distance is the distance to the leaves in the last row.
Let’s say we are now considering a nodes at a depth d. Remember the depth is 0-indexed, the last row is (h-1), and so the second last row is (h-2). For nodes which have the last row in the subtree, the maximum distance is given by (h-1) - d, and for nodes which have only the second last row in their subtree, the maximum distance is (h-2) - d.
Now we know how to get the maximum distance for nodes on any given level. Since there are only h rows, i.e. O(\log N) rows, looping through all the levels is a sufficiently efficient algorithm. All that is left is to find out, on each level, how many nodes contain the last row in their subtree and how many do not. Once that is done, we can easily loop through all the rows and just compute the maximum distance at that depth for nodes with the last row in their subtree, and nodes without. If either of these distances is equal to X, add the corresponding number of nodes at that depth to the answer.
Now, note that the last row is filled in from the left. A consequence of this is that, if a node at some depth has the last row in its subtree, then all the nodes to its left also have the last row in their subtrees. Thus, if we find the right-most node that has the last row in its subtree, and this is the k^{th} node in that row, then we know that there are k nodes at that depth having the last row in their subtree. If the depth is d, then the number of nodes not having the last row in the subtree is obviously 2^d - k, since 2^d is the total number of nodes at that depth.
Now to find this right-most node. At a depth d, we want to find the right-most node that has a node from the last row in its subtree. Now, realize that, correspondingly, this right-most node must have, in its subtree, the right-most node in the last row. Convince yourself of this fact before moving on. This is the same as saying that the right-most node at depth d, that we want, will be one of the ancestors of the right-most node in the last row. Thus we start at the right-most node of the last row, and keep moving up through its ancestors, and we know how many nodes are to this node’s left in the current row.
So an algorithm like this will work: Start from the right-most node in the last row. Keep moving up its parents. If we come to a depth d such that (h-1) - d = X, then add the number of nodes including and to the left of the current node to the answer. If we come to a depth such that (h-2) - d = X, add the number of nodes to the right of the current node to the answer.
Refer to editorialist’s solution for implementation details.