Data-Structure-Python/Lazy_Segment_Tree.py at main · deepanshmittal/Data-Structure-Python · 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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
# https://codeforces.com/contest/1108/submission/48856997
# https://codeforces.com/blog/entry/18051
# https://codeforces.com/blog/entry/15890
# https://codeforces.com/blog/entry/15729


from math import log2, ceil
from sys import stdin
from bisect import bisect_left as bl
from collections import defaultdict

input = stdin.readline
read = lambda: map(int, input().strip().split())


# Function To return left node index
def left(idx):
    return 2 * idx


# Function To return right node index
def right(idx):
    return 2 * idx + 1


class SegmentTree:
    from math import ceil, log2

    # Build Tree to store the sum of array in some range at each index of the tree
    def __init__(self, arr):
        self.size = len(arr)
        self.tree_size = 2 ** self.size.bit_length()
        self.tree = [0] * 2 * self.tree_size
        self.lazy = self.tree.copy()
        for idx in range(self.size):
            self.tree[self.tree_size + idx] = arr[idx]
        for idx in reversed(range(1, self.tree_size)):
            self.tree[idx] = self.tree[left(idx)] + self.tree[right(idx)]
        self.idx, self.left_node, self.right_node = [1, 0, self.size - 1]

    # Return Sum of element present in [range_left,range_right] of array both inclusive
    def range_sum(self, range_left, range_right):
        sum = 0
        range_left += self.tree_size
        range_right += self.tree_size
        self.lazy_update(range_left, range_right)
        while range_left <= range_right:
            if range_left & 1:
                self.node_update(range_left // 2)
                sum += self.tree[range_left]
                range_left += 1
            if not range_right & 1:
                self.node_update(range_right // 2)
                sum += self.tree[range_right]
                range_right -= 1
            range_right //= 2
            range_left //= 2
        return sum

    # Function to update the changes in parent node and pass the information to child node
    def node_update(self, index):
        for node in [left(index), right(index)]:
            if self.lazy[node]:
                self.tree[node] = self.lazy[node]
                self.lazy[node] = 0
                if node < self.tree_size:
                    self.lazy[left(node)] = self.lazy[right(node)] = self.tree[node] // 2

    # Lazily updates the stored information to calculate the new information correctly
    def lazy_update(self, left_index, right_index):
        left_update = []
        right_update = []
        left_index //= 2
        right_index //= 2
        while left_index:
            left_update.append(left_index)
            right_update.append(right_index)
            left_index //= 2
            right_index //= 2
        for lst_update in [left_update, right_update]:
            for index in reversed(lst_update):
                if self.lazy[index]:
                    self.tree[index] = self.lazy[index]
                    self.lazy[index] = 0
                    self.lazy[left(index)] = self.lazy[right(index)] = self.tree[index] // 2

    # Range Update to assign a new_value to every element in range [range_left, range_right] both inclusive
    def range_update(self, range_left, range_right, new_val):
        range_left += self.tree_size
        range_right += self.tree_size
        start_left, start_right = range_left, range_right
        # Lazily update the nodes encountered in path (range_left--> index 1) and (range_right --> index 1)
        self.lazy_update(range_left, range_right)
        while range_left <= range_right:
            if range_left & 1:
                self.tree[range_left] = new_val
                self.lazy[range_left] = 0
                if range_left < self.tree_size:
                    self.lazy[left(range_left)] = self.tree[range_left] // 2
                    self.lazy[right(range_left)] = self.tree[range_left] // 2
                range_left += 1

            if not range_right & 1:
                self.tree[range_right] = new_val
                self.lazy[range_right] = 0
                if range_right < self.tree_size:
                    self.lazy[left(range_right)] = self.tree[range_right] // 2
                    self.lazy[right(range_right)] = self.tree[range_right] // 2
                range_right -= 1
            new_val *= 2
            range_left //= 2
            range_right //= 2
        self.sum_update(start_left, start_right)

    # Update Sum information to the parents of newly updated nodes
    def sum_update(self, start_left, start_right):
        start_left = start_left // 2
        start_right = start_right // 2
        while start_left:
            self.node_update(start_left)
            self.node_update(start_right)
            self.tree[start_left] = self.tree[left(start_left)] + self.tree[right(start_left)]
            self.tree[start_right] = self.tree[left(start_right)] + self.tree[right(start_right)]
            start_left //= 2
            start_right //= 2


# lst = list(range(1, 11))
# st = SegmentTree(lst)
# print(st.tree)
# st.range_update(0, 9, 10)
# print(st.tree)
# print(st.lazy)
# st.range_update(5, 9, 20)
# print(st.tree)
# print(st.lazy)
# print(st.range_sum(3, 6))
# print(st.tree)
# print(st.lazy)