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.
Solution:
class TreeAncestor {
/*
dp state:
dp[i][j] = j-th node's 2^(i)th ancestor in the path
dp init:
dp[i][j] = dp[0][j] (first parent (2^0) of each node is given)
dp trans:
dp[i][j] = dp[i-1][ dp[i-1][j] ]
meaning: A(j, 2^i) = A( A(j, 2^i-1), 2^i-1)
To find the (2^i)-th ancestor of j, recursively find j-th node's 2^(i-1)th ancestor's 2^(i-1)th ancestor. (2^(i) = 2^(i-1) + 2^(i-1))
*/
int[][] dp;
int maxPow;
public TreeAncestor(int n, int[] parent) {
// log_base_2(n)
maxPow = (int) (Math.log(n) / Math.log(2)) + 1;
dp = new int[maxPow][n];
dp[0] = parent;
for (int i = 1; i < maxPow; i ++) {
for (int j = 0; j < n; j ++) {
int pre = dp[i - 1][j];
dp[i][j] = pre == -1 ? -1 : dp[i - 1][pre];
}
}
}
public int getKthAncestor(int node, int k) {
int maxPow = this.maxPow;
while (node > -1 && k > 0) {
if (k >= (1 << maxPow)) {
node = dp[maxPow][node];
k -= (1 << maxPow);
} else {
maxPow --;
}
}
return node;
}
}
/**
* Your TreeAncestor object will be instantiated and called as such:
* TreeAncestor obj = new TreeAncestor(n, parent);
* int param_1 = obj.getKthAncestor(node,k);
*/