Preview:
class SegmentTreeLazyPropagation:
    def __init__(self, A):
        self.A = A
        self.tree = [0] * (4 * len(A))  # Adjust the size as needed
        self.lazy = [0] * (4 * len(A))
        self.build(0, 0, len(A) - 1)

    def build(self, node, start, end):
        if start == end:
            self.tree[node] = self.A[start]
        else:
            mid = (start + end) // 2
            self.build(2 * node + 1, start, mid)
            self.build(2 * node + 2, mid + 1, end)
            self.tree[node] = self.tree[2 * node + 1] + self.tree[2 * node + 2]

    def update_range(self, l, r, val):
        self._update_range(0, 0, len(self.A) - 1, l, r, val)

    def _update_range(self, node, start, end, l, r, val):
        if self.lazy[node] != 0:
            self.tree[node] += (end - start + 1) * self.lazy[node]
            if start != end:
                self.lazy[2 * node + 1] += self.lazy[node]
                self.lazy[2 * node + 2] += self.lazy[node]
            self.lazy[node] = 0

        if r < start or end < l:
            return

        if l <= start <= end <= r:
            self.tree[node] += (end - start + 1) * val
            if start != end:
                self.lazy[2 * node + 1] += val
                self.lazy[2 * node + 2] += val
            return

        mid = (start + end) // 2
        self._update_range(2 * node + 1, start, mid, l, r, val)
        self._update_range(2 * node + 2, mid + 1, end, l, r, val)
        self.tree[node] = self.tree[2 * node + 1] + self.tree[2 * node + 2]

    def query(self, l, r):
        return self._query(0, 0, len(self.A) - 1, l, r)

    def _query(self, node, start, end, l, r):
        if self.lazy[node] != 0:
            self.tree[node] += (end - start + 1) * self.lazy[node]
            if start != end:
                self.lazy[2 * node + 1] += self.lazy[node]
                self.lazy[2 * node + 2] += self.lazy[node]
            self.lazy[node] = 0

        if r < start or end < l:
            return 0

        if l <= start <= end <= r:
            return self.tree[node]

        mid = (start + end) // 2
        p1 = self._query(2 * node + 1, start, mid, l, r)
        p2 = self._query(2 * node + 2, mid + 1, end, l, r)
        return p1 + p2

# Example Usage
A = [1, 2, 3, 4, 5]
seg_tree_lazy = SegmentTreeLazyPropagation(A)

# Range Update example
seg_tree_lazy.update_range(1, 3, 2)
result = seg_tree_lazy.query(1, 3)
print("Sum in range [1, 3] after update:", result)

# Query example
result = seg_tree_lazy.query(2, 4)
print("Sum in range [2, 4]:", result)
downloadDownload PNG downloadDownload JPEG downloadDownload SVG

Tip: You can change the style, width & colours of the snippet with the inspect tool before clicking Download!

Click to optimize width for Twitter