Kth Ancestor of a Tree Node

hard tree binary lifting

Problem

You are given a tree with n nodes numbered 0 to n − 1, where node 0 is the root and parent[i] is the parent of node i (parent[0] = −1). Design a data structure answering queries getKthAncestor(node, k): the node reached after walking k steps up toward the root, or −1 if node has fewer than k ancestors.

With up to 5·10⁴ nodes and 5·10⁴ queries (1 ≤ k ≤ n), and trees that can be long chains, walking up one parent at a time is too slow — each query needs to run in roughly logarithmic time.

Inputparent = [-1,0,0,1,1,2,2]; getKthAncestor(3,1), getKthAncestor(5,2), getKthAncestor(6,3)

Output1, 0, -1

One step above 3 is 1, two steps above 5 is the root 0, and 6 has only two ancestors so a third does not exist.

parent (comma-separated, parent[0] = -1)

queries (node:k, comma-separated)

class TreeAncestor:
    def __init__(self, n, parent):
        self.LOG = n.bit_length()
        up = [parent[:]]                  # level 0 = direct parents
        for j in range(1, self.LOG):
            prev = up[j - 1]
            up.append([-1 if prev[v] == -1 else prev[prev[v]]
                       for v in range(n)])
        self.up = up

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

class TreeAncestor {
  constructor(n, parent) {
    this.LOG = 32 - Math.clz32(n);
    this.up = [parent.slice()];          // level 0 = direct parents
    for (let j = 1; j < this.LOG; j++) {
      const prev = this.up[j - 1], row = [];
      for (let v = 0; v < n; v++)
        row.push(prev[v] === -1 ? -1 : prev[prev[v]]);
      this.up.push(row);
    }
  }

  getKthAncestor(node, k) {
    for (let j = 0; j < this.LOG; j++) {
      if (k >> j & 1) {
        node = this.up[j][node];
        if (node === -1) return -1;
      }
    }
    return node;
  }
}

class TreeAncestor {
    private final int LOG;
    private final int[][] up;

    public TreeAncestor(int n, int[] parent) {
        LOG = 32 - Integer.numberOfLeadingZeros(n);
        up = new int[LOG][n];
        up[0] = parent.clone();          // level 0 = direct parents
        for (int j = 1; j < LOG; j++)
            for (int v = 0; v < n; v++)
                up[j][v] = up[j-1][v] == -1
                    ? -1 : up[j-1][up[j-1][v]];
    }

    public int getKthAncestor(int node, int k) {
        for (int j = 0; j < LOG; j++) {
            if ((k >> j & 1) == 1) {
                node = up[j][node];
                if (node == -1) return -1;
            }
        }
        return node;
    }
}

class TreeAncestor {
    int LOG;
    vector<vector<int>> up;
public:
    TreeAncestor(int n, vector<int>& parent) {
        LOG = 32 - __builtin_clz(n);
        up.assign(LOG, vector<int>(n, -1));
        up[0] = parent;                  // level 0 = direct parents
        for (int j = 1; j < LOG; j++)
            for (int v = 0; v < n; v++)
                if (up[j-1][v] != -1)
                    up[j][v] = up[j-1][up[j-1][v]];
    }

    int getKthAncestor(int node, int k) {
        for (int j = 0; j < LOG; j++) {
            if (k >> j & 1) {
                node = up[j][node];
                if (node == -1) return -1;
            }
        }
        return node;
    }
};

Explanation

The naive answer walks up one parent at a time, costing O(k) per query — on a chain of 5·10⁴ nodes with 5·10⁴ queries that is billions of steps. Binary lifting fixes this by precomputing jumps of power-of-two sizes.

We build a table up[j][v] = the 2^j-th ancestor of v (or −1 if it does not exist). Level 0 is just the parent array. Each higher level composes two jumps of the level below: up[j][v] = up[j-1][up[j-1][v]] — jumping 2^(j-1) steps twice lands exactly 2^j steps up. With LOG ≈ log₂(n) levels the table covers every jump size we could need.

A query then writes k in binary: every number is a sum of distinct powers of two, so walking k steps decomposes into one table lookup per set bit. For each bit j of k we replace node with up[j][node]; if we ever hit −1, fewer than k ancestors exist and the answer is −1.

Default example: parent = [-1,0,0,1,1,2,2], so LOG = 3. For getKthAncestor(5, 2), k = 10₂ needs only the bit-1 jump: up[1][5] = 0, the root. For getKthAncestor(6, 3), k = 11₂ first jumps 1 step to node 2, then asks for its 2-step ancestor: up[1][2] = -1, so the answer is −1.

Building the table costs O(n log n) time and space, and each query touches at most log₂(k) + 1 cells, so queries run in O(log n) — fast enough for any mix of deep chains and heavy query loads.

Time: O(n log n) build, O(log n) per query Space: O(n log n)