Kth Ancestor of a Tree Node

Sanjeev SharmaSanjeev Sharma
3 min read

Advertisement

Problem

Given a tree (not binary tree), implement getKthAncestor(node, k) returning the kth ancestor or -1 if none.

Key insight: Binary lifting. Precompute ancestor[node][j] = 2^j-th ancestor. Answer k in O(log k) by decomposing k in binary.

Approach — Binary Lifting

ancestor[v][0] = parent[v]
ancestor[v][j] = ancestor[ancestor[v][j-1]][j-1]

Solutions

// C++
class TreeAncestor {
    int LOG = 16;
    vector<vector<int>> up;
public:
    TreeAncestor(int n, vector<int>& parent) {
        up.assign(n, vector<int>(LOG, -1));
        for (int i = 0; i < n; i++) up[i][0] = parent[i];
        for (int j = 1; j < LOG; j++)
            for (int i = 0; i < n; i++)
                if (up[i][j-1] != -1) up[i][j] = up[up[i][j-1]][j-1];
    }
    int getKthAncestor(int node, int k) {
        for (int j = 0; j < LOG; j++)
            if ((k >> j) & 1) {
                node = up[node][j];
                if (node == -1) return -1;
            }
        return node;
    }
};
// Java
class TreeAncestor {
    int LOG = 16;
    int[][] up;
    public TreeAncestor(int n, int[] parent) {
        up = new int[n][LOG];
        for (int[] row : up) Arrays.fill(row, -1);
        for (int i = 0; i < n; i++) up[i][0] = parent[i];
        for (int j = 1; j < LOG; j++)
            for (int i = 0; i < n; i++)
                if (up[i][j-1] != -1) up[i][j] = up[up[i][j-1]][j-1];
    }
    public int getKthAncestor(int node, int k) {
        for (int j = 0; j < LOG; j++)
            if (((k >> j) & 1) == 1) {
                node = up[node][j];
                if (node == -1) return -1;
            }
        return node;
    }
}
// JavaScript
class TreeAncestor {
    constructor(n, parent) {
        this.LOG = 16;
        this.up = Array.from({length: n}, () => Array(this.LOG).fill(-1));
        for (let i = 0; i < n; i++) this.up[i][0] = parent[i];
        for (let j = 1; j < this.LOG; j++)
            for (let i = 0; i < n; i++)
                if (this.up[i][j-1] !== -1)
                    this.up[i][j] = this.up[this.up[i][j-1]][j-1];
    }
    getKthAncestor(node, k) {
        for (let j = 0; j < this.LOG; j++)
            if ((k >> j) & 1) {
                node = this.up[node][j];
                if (node === -1) return -1;
            }
        return node;
    }
}
# Python
class TreeAncestor:
    def __init__(self, n, parent):
        LOG = 16
        self.up = [[-1] * LOG for _ in range(n)]
        for i in range(n):
            self.up[i][0] = parent[i]
        for j in range(1, LOG):
            for i in range(n):
                if self.up[i][j-1] != -1:
                    self.up[i][j] = self.up[self.up[i][j-1]][j-1]

    def getKthAncestor(self, node, k):
        for j in range(16):
            if (k >> j) & 1:
                node = self.up[node][j]
                if node == -1:
                    return -1
        return node

Complexity

  • Preprocessing: O(n log n)
  • Query: O(log k)
  • Space: O(n log n)

Key Insight

Binary lifting = "doubling trick." Jump by powers of 2, decompose k in binary. Used in LCA algorithms too.

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro