|
| 1 | +import java.util.LinkedList; |
| 2 | +/** |
| 3 | +* Time: O(N log N + Q * log N), each query is answered in log N time. Space: O(N log N) |
| 4 | +* Use: |
| 5 | +* Your BinaryLifting object will be instantiated and called as such: |
| 6 | +* BinaryLifting obj = new BinaryLifting(n, parent); |
| 7 | +* int param_1 = obj.getKthAncestor(node,k); |
| 8 | +* ref: https://leetcode.com/problems/kth-ancestor-of-a-tree-node/ and https://www.youtube.com/watch?v=oib-XsjFa-M |
| 9 | +*/ |
| 10 | +class BinaryLifting { |
| 11 | + // preprocess |
| 12 | + // O(N log N) |
| 13 | + // precompute the answer for power of 2 |
| 14 | + private int[][] atLevel; // atLevel[nodeId][level] means what is the predecessor at 2^level higher |
| 15 | + private int MAX_LOG = 0; |
| 16 | + boolean vis[]; |
| 17 | + public BinaryLifting(int n, int[] parent) { |
| 18 | + MAX_LOG = 0; |
| 19 | + vis = new boolean[n]; |
| 20 | + while(n >= (1 << MAX_LOG)){ |
| 21 | + MAX_LOG++; |
| 22 | + } |
| 23 | + atLevel = new int[n][MAX_LOG]; |
| 24 | + for(int nodeId = 0; nodeId < n; nodeId++){ |
| 25 | + for(int level = 0; level < MAX_LOG; level++){ |
| 26 | + atLevel[nodeId][level] = -1; |
| 27 | + } |
| 28 | + } |
| 29 | + for(int nodeId = 1; nodeId <= n - 1; nodeId++){ |
| 30 | + if(vis[nodeId])continue; |
| 31 | + LinkedList<Integer> unVisited = new LinkedList<Integer>(); // linked list as a stack for unvisited node |
| 32 | + int currentNode = nodeId; |
| 33 | + while(currentNode != -1 && !vis[currentNode]){ |
| 34 | + unVisited.addLast(currentNode); |
| 35 | + currentNode = parent[currentNode]; |
| 36 | + } |
| 37 | + while(!unVisited.isEmpty()){ |
| 38 | + int topUnvisitedNode = unVisited.removeLast(); |
| 39 | + atLevel[topUnvisitedNode][0] = parent[topUnvisitedNode]; |
| 40 | + for(int level = 1; level <= MAX_LOG - 1; level++){ |
| 41 | + if(atLevel[topUnvisitedNode][level - 1] != -1){ |
| 42 | + atLevel[topUnvisitedNode][level] = atLevel[atLevel[topUnvisitedNode][level - 1]][level - 1]; |
| 43 | + }else{ |
| 44 | + break; |
| 45 | + } |
| 46 | + } |
| 47 | + vis[topUnvisitedNode] = true; |
| 48 | + } |
| 49 | + } |
| 50 | + } |
| 51 | + |
| 52 | + public int getKthAncestor(int node, int k) { |
| 53 | + int kthAncestor = node; |
| 54 | + for(int level = MAX_LOG - 1; level >= 0; level--){ // at ancestor at 2^level |
| 55 | + if((k & (1 << level)) > 0){ // check if ith bit is set |
| 56 | + // every numer can be represented by sum of power of 2 |
| 57 | + kthAncestor = atLevel[kthAncestor][level]; |
| 58 | + if(kthAncestor == -1){ |
| 59 | + break; |
| 60 | + } |
| 61 | + } |
| 62 | + } |
| 63 | + return kthAncestor; |
| 64 | + } |
| 65 | +} |
| 66 | + |
0 commit comments