python-lca/lca.py at main · 6851-2021/python-lca · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
from pm_rmq import pm_rmq_query, pm_rmq_preprocess

class Node:
    def __init__(self, data):
        self.left = None
        self.right = None
        self.first_visit = None
        self.data = data

### Least Common Ancestor ###

def lca_preprocessing(root):
    depth_arr, node_arr = binary_tree_to_arr(root)
    preprocessed_pm_arr = pm_rmq_preprocess(depth_arr)
    return (preprocessed_pm_arr, depth_arr, node_arr)

def lca_query(node1, node2, preprocessed_tree):
    preprocessed_pm_arr, depth_arr, node_arr = preprocessed_tree

    start_index = min(node1.first_visit, node2.first_visit)
    end_index = max(node1.first_visit, node2.first_visit)
    min_index = pm_rmq_query(preprocessed_pm_arr, start_index, end_index)
    return node_arr[min_index]

### Helper Methods ###

def euler_tour(curr_node, curr_depth, depth_arr, node_arr):
    if curr_node is None:
        return
    else:
        if curr_node.first_visit is None:
            curr_node.first_visit = len(depth_arr)
        depth_arr.append(curr_depth)
        node_arr.append(curr_node)
        if curr_node.left is not None:
            euler_tour(curr_node.left, curr_depth + 1, depth_arr, node_arr)
            depth_arr.append(curr_depth)
            node_arr.append(curr_node)
        if curr_node.right is not None:
            euler_tour(curr_node.right, curr_depth + 1, depth_arr, node_arr)
            depth_arr.append(curr_depth)
            node_arr.append(curr_node)

def binary_tree_to_arr(root):
    depth_arr = []
    node_arr = []
    euler_tour(root, 0, depth_arr, node_arr)
    return depth_arr, node_arr

if __name__ == '__main__':
    # Example tree from L15
    a = Node('0-root')
    b = Node('1-left')
    c = Node('2-left-left')
    d = Node('1-right')
    e = Node('2-right-left')
    f = Node('3-left')
    g = Node('3-right')
    h = Node('2-right-right')
    a.left = b
    b.left = c
    a.right = d
    d.left = e
    e.left = f
    e.right = g
    d.right = h

    preprocessed_tree = lca_preprocessing(a)
    node = lca_query(h, c, preprocessed_tree)
    print("Example:")
    print("LCA of {} and {}".format(h.data, c.data))
    print("    {}".format(node.data))