LeetCode Solutions Blog

Kth Ancestor of a Tree Node

Number: 1296

Difficulty: Hard

Paid? No

Problem Description

Given a tree with n nodes (from 0 to n - 1) represented by a parent array, implement a class that can quickly return the kth ancestor of any given node. The kth ancestor is the kth node in the path from the node to the root. If no such ancestor exists, return -1.

Key Insights

Precompute ancestors using Binary Lifting to answer kth ancestor queries in O(log n) time.
Use a 2D array (or table) where dp[node][j] represents the 2^j-th ancestor of the node.
Precompute the table for all nodes in O(n log n) time.
Each query is then answered by decomposing k into powers of two and lifting the node up accordingly.
This approach is efficient given the constraints (up to 50000 nodes and queries).

Space and Time Complexity

Time Complexity:

Preprocessing: O(n log n)
Each query: O(log n) Space Complexity: O(n log n) for storing the binary lifting table.

Solution

The solution utilizes binary lifting, where we precompute for every node the ancestor at powers of two distances (i.e. 1, 2, 4, 8, ...). During a query for the kth ancestor of a node, we iterate through the bits in the binary representation of k. For each set bit, we move the node pointer to its corresponding ancestor using our precomputed table. If at any point the node becomes -1, we return -1 indicating that the kth ancestor does not exist.

Code Solutions

class TreeAncestor:
    def __init__(self, n: int, parent: list):
        # Calculate max power for binary lifting table
        self.max_power = 1
        while (1 << self.max_power) <= n:
            self.max_power += 1

        # Initialize dp table with dimensions n x max_power, default to -1
        self.dp = [[-1] * self.max_power for _ in range(n)]
        
        # The first ancestor (2^0) is the direct parent
        for i in range(n):
            self.dp[i][0] = parent[i]
        
        # Fill the table: dp[i][j] is the 2^j-th ancestor of node i
        for j in range(1, self.max_power):
            for i in range(n):
                prev_ancestor = self.dp[i][j - 1]
                if prev_ancestor != -1:
                    self.dp[i][j] = self.dp[prev_ancestor][j - 1]
                    
    def getKthAncestor(self, node: int, k: int) -> int:
        # Process bits of k: lift node by 2^j when corresponding bit is set
        for j in range(self.max_power):
            if k & (1 << j):
                node = self.dp[node][j]
                if node == -1:
                    return -1
        return node

# Example usage:
# treeAncestor = TreeAncestor(7, [-1, 0, 0, 1, 1, 2, 2])
# print(treeAncestor.getKthAncestor(3, 1))  # Output: 1
# print(treeAncestor.getKthAncestor(5, 2))  # Output: 0
# print(treeAncestor.getKthAncestor(6, 3))  # Output: -1