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)
Preview:
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