|
| 1 | +import queue |
| 2 | + |
| 3 | + |
| 4 | +def swap(a, b): |
| 5 | + a ^= b |
| 6 | + b ^= a |
| 7 | + a ^= b |
| 8 | + return a, b |
| 9 | + |
| 10 | + |
| 11 | +# creating sparse table which saves each nodes 2^ith parent |
| 12 | +def creatSparse(max_node, parent): |
| 13 | + j = 1 |
| 14 | + while (1 << j) < max_node: |
| 15 | + for i in range(1, max_node + 1): |
| 16 | + parent[j][i] = parent[j - 1][parent[j - 1][i]] |
| 17 | + j += 1 |
| 18 | + return parent |
| 19 | + |
| 20 | + |
| 21 | +# returns lca of node u,v |
| 22 | +def LCA(u, v, level, parent): |
| 23 | + # u must be deeper in the tree than v |
| 24 | + if level[u] < level[v]: |
| 25 | + u, v = swap(u, v) |
| 26 | + # making depth of u same as depth of v |
| 27 | + for i in range(18, -1, -1): |
| 28 | + if level[u] - (1 << i) >= level[v]: |
| 29 | + u = parent[i][u] |
| 30 | + # at the same depth if u==v that mean lca is found |
| 31 | + if u == v: |
| 32 | + return u |
| 33 | + # moving both nodes upwards till lca in found |
| 34 | + for i in range(18, -1, -1): |
| 35 | + if parent[i][u] != 0 and parent[i][u] != parent[i][v]: |
| 36 | + u, v = parent[i][u], parent[i][v] |
| 37 | + # returning longest common ancestor of u,v |
| 38 | + return parent[0][u] |
| 39 | + |
| 40 | + |
| 41 | +# runs a breadth first search from root node of the tree |
| 42 | +# sets every nodes direct parent |
| 43 | +# parent of root node is set to 0 |
| 44 | +# calculates depth of each node from root node |
| 45 | +def bfs(level, parent, max_node, graph, root=1): |
| 46 | + level[root] = 0 |
| 47 | + q = queue.Queue(maxsize=max_node) |
| 48 | + q.put(root) |
| 49 | + while q.qsize() != 0: |
| 50 | + u = q.get() |
| 51 | + for v in graph[u]: |
| 52 | + if level[v] == -1: |
| 53 | + level[v] = level[u] + 1 |
| 54 | + q.put(v) |
| 55 | + parent[0][v] = u |
| 56 | + return level, parent |
| 57 | + |
| 58 | + |
| 59 | +def main(): |
| 60 | + max_node = 13 |
| 61 | + # initializing with 0 |
| 62 | + parent = [[0 for _ in range(max_node + 10)] for _ in range(20)] |
| 63 | + # initializing with -1 which means every node is unvisited |
| 64 | + level = [-1 for _ in range(max_node + 10)] |
| 65 | + graph = { |
| 66 | + 1: [2, 3, 4], |
| 67 | + 2: [5], |
| 68 | + 3: [6, 7], |
| 69 | + 4: [8], |
| 70 | + 5: [9, 10], |
| 71 | + 6: [11], |
| 72 | + 7: [], |
| 73 | + 8: [12, 13], |
| 74 | + 9: [], |
| 75 | + 10: [], |
| 76 | + 11: [], |
| 77 | + 12: [], |
| 78 | + 13: [] |
| 79 | + } |
| 80 | + level, parent = bfs(level, parent, max_node, graph, 1) |
| 81 | + parent = creatSparse(max_node, parent) |
| 82 | + print("LCA of node 1 and 3 is: ", LCA(1, 3, level, parent)) |
| 83 | + print("LCA of node 5 and 6 is: ", LCA(5, 6, level, parent)) |
| 84 | + print("LCA of node 7 and 11 is: ", LCA(7, 11, level, parent)) |
| 85 | + print("LCA of node 6 and 7 is: ", LCA(6, 7, level, parent)) |
| 86 | + print("LCA of node 4 and 12 is: ", LCA(4, 12, level, parent)) |
| 87 | + print("LCA of node 8 and 8 is: ", LCA(8, 8, level, parent)) |
| 88 | + |
| 89 | + |
| 90 | +if __name__ == "__main__": |
| 91 | + main() |
0 commit comments