LC 1483 [H] Kth Ancestor of a Tree Node - ALawliet/algorithms GitHub Wiki
Algo For node 0, 1, ..., n-1, we define a matrix self.dp[][] whose i, jth element indicates the ith node's 2^j parent. Here, i = 0, 1, ..., n-1 and j = 0, 1, ..., int(log2(n)). An important recursive relationship is that
self.dp[i][j] = self.dp[self.dp[i][j-1]][j-1].
In other words, ith node's 2^j parent is ith node's 2^(j-1) parent's 2^(j-1) parent. In this way, lookup is guarenteed to complete in O(logN) time. Note that it takes O(NlogN) to build self.dp.
Time complexity O(NlogN) to pre-processing the tree and O(logN) for each equery thereafter.
class TreeAncestor:
def __init__(self, n: int, parent: List[int]):
LOG = int(log2(n)) + 1 #at most 16 for this problem, 2^x = n = LOG
self.up = [[-1] * LOG for _ in range(n)] #ith node's 2^j parent = up
# dp[num][pow2]
for j in range(LOG): # pow2 = LOG
for v in range(n): # num = v
if j == 0:
self.up[v][0] = parent[v] #2^0 parent
elif self.up[v][j-1] != -1: # binary jumps
self.up[v][j] = self.up[ self.up[v][j-1] ][j-1]
# parent[0] = 0
# for v in range(n):
# self.dp[v][0] = parent[v]
# for j in range(1, m):
# self.dp[v][j] = self.dp[ self.dp[v][j-1] ][j-1]
def getKthAncestor(self, node: int, k: int) -> int:
while k > 0 and node != -1:
i = int(log2(k))
node = self.up[node][i]
k -= (1 << i) # jth bit is ON in number k
return node
'''
k = 19 = 10011 = 16 + 2 + 1
int up[N][LOG]
up[v][k] -- 2^j-th ancestor of v (j=5 is 2^5=32nd ancestor of v)
for v = 0 .. N-1:
up[v][0] = parent[v] # kth ancestor=1, up by 1
---
up[v][1] = up[ up[v][0] ][0] # k=2, grandparent (parent of parent)
up[v][2] = up[ up[v][1] ][1] # k=4, (grandparent of grandparent = 4th ancestor = 2nd ancestor of 2nd ancestor)
up[v][3] = up[ up[v][2] ][2] # k=8, (up by 4 and up by 4)
for j = 1 .. LOG - 1:
up[v][j] = up[ up[v][j-1] ][j-1]
O(N*log(N))
''' // Binary Lifting tutorial https://youtu.be/oib-XsjFa-M
class TreeAncestor {
vector<vector<int>> up; // int up[N][20];
vector<int> depth;
int LOG;
public:
TreeAncestor(int n, vector<int>& parent) {
LOG = 0;
while((1 << LOG) <= n) {
LOG++;
}
up = vector<vector<int>>(n, vector<int>(LOG));
depth = vector<int>(n);
// up[v][j] is 2^j -th ancestor of node v
parent[0] = 0;
for(int v = 0; v < n; v++) {
up[v][0] = parent[v];
if(v != 0) {
depth[v] = depth[parent[v]] + 1;
}
for(int j = 1; j < LOG; j++) {
up[v][j] = up[ up[v][j-1] ][j-1];
}
}
}
int getKthAncestor(int node, int k) {
if(depth[node] < k) {
return -1;
}
for(int j = LOG - 1; j >= 0; j--) {
if(k >= (1 << j)) {
node = up[node][j];
k -= 1 << j;
}
}
return node;
}
};