Program to find Kth ancestor of a tree node in Python

Suppose we have a tree with n nodes that are numbered from 0 to n-1. The tree is given by a parent array, where parent[i] is the parent of node i. The root of the tree is node 0. We have to find the kth ancestor of a given node. If the ancestor is not present, then return -1.

So, if the input is like ?

Then the output will be 2 because the first ancestor of node 6 is 5 and the second is 2.

Algorithm

To solve this, we will follow these steps ?

• Define a function solve() that takes parent, node, k

• If node is same as -1, then return -1

• If k is same as 1, then return parent[node]

• If (k AND k-1) is zero (k is power of 2), then recursively find k/2th ancestor twice

• Otherwise, find the most significant bit (MSB) and split the problem

Example

Let us see the following implementation to get better understanding ?

def solve(parent, node, k):
    if node == -1:
        return -1
    elif k == 1:
        return parent[node]
    elif not (k & k-1):
        return solve(parent, solve(parent, node, k >> 1), k >> 1)
    else:
        msb = 1 << (k.bit_length()-1)
        return solve(parent, solve(parent, node, k-msb), msb)

parent = [-1, 0, 0, 1, 2, 2, 5, 5]
node = 6
k = 2
print(solve(parent, node, k))

2

How It Works

The algorithm uses binary lifting technique to efficiently find the kth ancestor:

• Base case: If node is -1 (no parent), return -1

• Direct case: If k=1, return the immediate parent

• Power of 2: If k is a power of 2, find k/2th ancestor twice

• General case: Split k using its most significant bit

Alternative Approach

A simpler iterative approach for finding kth ancestor ?

def find_kth_ancestor(parent, node, k):
    current = node
    for i in range(k):
        if current == -1:
            return -1
        current = parent[current]
    return current

# Test with the same example
parent = [-1, 0, 0, 1, 2, 2, 5, 5]
node = 6
k = 2
result = find_kth_ancestor(parent, node, k)
print(f"The {k}th ancestor of node {node} is: {result}")

The 2nd ancestor of node 6 is: 2

Comparison

Conclusion

The binary lifting approach efficiently finds the kth ancestor in O(log k) time using recursive calls. For simple cases, the iterative approach is more straightforward and uses constant space.