Formatted question description: https://leetcode.ca/all/1483.html

1483. Kth Ancestor of a Tree Node (Hard)

You are given a tree with n nodes numbered from 0 to n-1 in the form of a parent array where parent[i] is the parent of node i. The root of the tree is node 0.

Implement the function getKthAncestor(int node, int k) to return the k-th ancestor of the given node. If there is no such ancestor, return -1.

The k-th ancestor of a tree node is the k-th node in the path from that node to the root.

Example:

Input:
["TreeAncestor","getKthAncestor","getKthAncestor","getKthAncestor"]
[[7,[-1,0,0,1,1,2,2]],[3,1],[5,2],[6,3]]

Output:
[null,1,0,-1]

Explanation:
TreeAncestor treeAncestor = new TreeAncestor(7, [-1, 0, 0, 1, 1, 2, 2]);

treeAncestor.getKthAncestor(3, 1);  // returns 1 which is the parent of 3
treeAncestor.getKthAncestor(5, 2);  // returns 0 which is the grandparent of 5
treeAncestor.getKthAncestor(6, 3);  // returns -1 because there is no such ancestor

Constraints:

1 <= k <= n <= 5*10^4
parent[0] == -1 indicating that 0 is the root node.
0 <= parent[i] < n for all 0 < i < n
0 <= node < n
There will be at most 5*10^4 queries.

Related Topics:
Dynamic Programming

Solution 1.

Assign increasing ID to nodes via preorder traversal, then you will get the following relationship

      0                       0
    /   \                   /   \
   1     2       =>        1     4
  / \   / \               / \   / \
 3   4 5   6             2   3 5   6

We store the tree of the id’s using vector<vector<int>>

[
    [0],
    [1,4],
    [2,3,5,6]
]

Then given a node, we can find the corresponding id[node]. For example id[2] = 4, id[5] = 5.

Also we can store the corresponding level of node. For example level[5] = 2.

Then if we want to find getKthAncestor(5, 1), we can know that its parent is at level level[5] - 1 = 1.

Then in level 1, we use binary search to find the largest id that is smaller than id[5] = 5. The parent id is 4.

Then we use the id to find the corresponding node, 2.

// OJ: https://leetcode.com/problems/kth-ancestor-of-a-tree-node/

// Time:
//      TreeAncestor: O(N)
//      getKthAncestor: O(logN)
// Space: O(N)
class TreeAncestor {
    vector<int> nodeToId, idToNode, idToLevel;
    vector<vector<int>> tree, G;
    void preorder(int node, int &id, int level) {
        if (level >= tree.size()) tree.emplace_back();
        tree[level].push_back(id);
        nodeToId[node] = id;
        idToNode[id] = node;
        idToLevel[id] = level;
        ++id;
        for (int child : G[node]) preorder(child, id, level + 1);
    }
public:
    TreeAncestor(int n, vector<int>& parent) {
        nodeToId.assign(n, 0);
        idToNode.assign(n, 0);
        idToLevel.assign(n, 0);
        G.assign(n, {});
        for (int i = 1; i < n; ++i) G[parent[i]].push_back(i);
        int id = 0;
        preorder(0, id, 0);
    }
    int getKthAncestor(int node, int k) {
        int id = nodeToId[node], level = idToLevel[id] - k;
        if (level < 0) return -1;
        return idToNode[*(upper_bound(begin(tree[level]), end(tree[level]), id) - 1)];
    }
};

Solution 2. Binary Lifting

The brute force solution of this problem requires O(K) time for querying, which will result in TLE.

If we precompute all the kth ancester of node, the query will be O(1), but the preparation will take O(K^2).

One solution in between which doesn’t require storing that much of data and also can boost the efficiency, is binary lifting.

The idea is that we break down the value based on binary representation.

So instead of storeing all 1, 2, 3, 4, 5, ... values, we only store 1, 2, 4, 8, ... values. We can use those 1, 2, 4, 8, ... values to restore those values in between.

For example, a 6th (110 in base 2) parent is the same as the 4th (100 in base 2) parent of the 2nd (010 in base 2) parent.

Let P[i][node] be the node’s 2^ith parent.

P[i][node] = P[i-1][ P[i-1][node] ]

The 6th parent of node can be expressed as P[2][ P[1][node] ].

// OJ: https://leetcode.com/problems/kth-ancestor-of-a-tree-node/

// Time:
//      TreeAncestor: O(log(max(K)) * N)
//      getKthAncestor: O(logK)
// Space: O(N)
// Ref: https://leetcode.com/problems/kth-ancestor-of-a-tree-node/discuss/686268/Explanation-for-this-question-c%2B%2B-sample-code
class TreeAncestor {
    vector<vector<int>> P;
public:
    TreeAncestor(int n, vector<int>& parent) {
        P.assign(20, vector<int>(parent.size(), -1));
        for (int i = 0; i < parent.size(); ++i) P[0][i] = parent[i]; // 2^0
        for (int i = 1; i < 20; ++i) { // 2^i
            for (int node = 0; node < parent.size(); ++node) {
                int p = P[i - 1][node];
                if (p != -1) P[i][node] = P[i - 1][p]; // P[i][node] = P[i - 1][ P[i - 1][node] ]
            }
        }
    }
    int getKthAncestor(int node, int k) {
        for (int i = 0; i < 20; ++i) {
            if ((k & (1 << i)) == 0) continue; // If `k` has `0` at `i`th bit, skip.
            node = P[i][node];
            if (node == -1) return -1;
        }
        return node;
    }
};

1483 - Kth Ancestor of a Tree Node

1483. Kth Ancestor of a Tree Node (Hard)

Solution 1.

Solution 2. Binary Lifting

All Problems

All Solutions