class LazySegmentTreeWithSetUpdate:
def __init__(self, v):
# Constructor: Initializes the segment tree parameters
self.n = 0
self.MYSIZE = 0
self.tree = []
self.lazy = []
self.stamp = []
self.pending = []
self.arr = []
self.tm = 0
self.init(len(v))
# Copy the input array to the internal array
for i in range(len(v)):
self.arr[i] = v[i]
# Post initialization to construct the segment tree
self.post_init()
def init(self, n):
# Initializes the segment tree parameters
self.n = n
self.MYSIZE = 4 * n + 10
# Initialize arrays for the segment tree
self.tree = [0] * self.MYSIZE
self.lazy = [0] * self.MYSIZE
self.arr = [0] * self.MYSIZE
self.stamp = [0] * self.MYSIZE
self.pending = [False] * self.MYSIZE
def post_init(self):
# Constructs the segment tree
self.constructST()
def updateRangeUtil(self, si, ss, se, us, ue, diff, stamp_value):
# Utility function to update a range lazily
if self.pending[si]:
# Update the tree value if there is a pending update
self.tree[si] = (se - ss + 1) * self.lazy[si]
if ss != se:
# Propagate the update to the children
if self.stamp[si * 2 + 1] <= stamp_value:
self.lazy[si * 2 + 1] = self.lazy[si]
self.stamp[si * 2 + 1] = stamp_value
self.pending[si * 2 + 1] = True
if self.stamp[si * 2 + 2] <= stamp_value:
self.lazy[si * 2 + 2] = self.lazy[si]
self.stamp[si * 2 + 2] = stamp_value
self.pending[si * 2 + 2] = True
self.pending[si] = False
# Check for out of bounds or no overlap
if ss > se or ss > ue or se < us:
return
# Total overlap, update the current node
if ss >= us and se <= ue:
self.tree[si] = (se - ss + 1) * diff
if ss != se:
# Propagate the update to the children
if self.stamp[si * 2 + 1] <= stamp_value:
self.lazy[si * 2 + 1] = diff
self.stamp[si * 2 + 1] = stamp_value
self.pending[si * 2 + 1] = True
if self.stamp[si * 2 + 2] <= stamp_value:
self.lazy[si * 2 + 2] = diff
self.stamp[si * 2 + 2] = stamp_value
self.pending[si * 2 + 2] = True
self.stamp[si] = stamp_value
return
# Partial overlap, update children
mid = (ss + se) // 2
self.updateRangeUtil(2 * si + 1, ss, mid, us, ue, diff, stamp_value)
self.updateRangeUtil(2 * si + 2, mid + 1, se, us, ue, diff, stamp_value)
# Update the current node based on children
self.tree[si] = self.tree[2 * si + 1] + self.tree[2 * si + 2]
def updateRange(self, us, ue, diff):
# Update a range in the segment tree
self.tm += 1
self.updateRangeUtil(0, 0, self.n - 1, us, ue, diff, self.tm)
def update(self, us, ue, diff):
# Update function, same as updateRange
self.updateRange(us, ue, diff)
def getSumUtil(self, ss, se, qs, qe, si):
# Utility function to get the sum in a range
if self.pending[si]:
# Update the tree value if there is a pending update
self.tree[si] = (se - ss + 1) * self.lazy[si]
if ss != se:
# Propagate the update to the children
if self.stamp[si * 2 + 1] <= self.stamp[si]:
self.lazy[si * 2 + 1] = self.lazy[si]
self.stamp[si * 2 + 1] = self.stamp[si]
self.pending[si * 2 + 1] = True
if self.stamp[si * 2 + 2] <= self.stamp[si]:
self.lazy[si * 2 + 2] = self.lazy[si]
self.stamp[si * 2 + 2] = self.stamp[si]
self.pending[si * 2 + 2] = True
self.pending[si] = False
# Check for out of bounds or no overlap
if ss > se or ss > qe or se < qs:
return 0
# Total overlap, return the value of the current node
if ss >= qs and se <= qe:
return self.tree[si]
# Partial overlap, get sum from children
mid = (ss + se) // 2
return self.getSumUtil(ss, mid, qs, qe, 2 * si + 1) + self.getSumUtil(mid + 1, se, qs, qe, 2 * si + 2)
def getSum(self, qs, qe):
# Get the sum in a range
if qs < 0 or qe > self.n - 1 or qs > qe:
print("Invalid Input")
return -1
return self.getSumUtil(0, self.n - 1, qs, qe, 0)
def query(self, qs, qe):
# Query function, same as getSum
return self.getSum(qs, qe)
def constructSTUtil(self, ss, se, si):
# Utility function to construct the segment tree
if ss > se:
return
if ss == se:
self.tree[si] = self.arr[ss]
return
mid = (ss + se) // 2
if ss <= mid:
self.constructSTUtil(ss, mid, 2 * si + 1)
if mid + 1 <= se:
self.constructSTUtil(mid + 1, se, 2 * si + 2)
# Update the current node based on children
self.tree[si] = self.tree[2 * si + 1] + self.tree[2 * si + 2]
def constructST(self):
# Construct the segment tree
self.constructSTUtil(0, self.n - 1, 0)
@staticmethod
def how_to_use():
# Example usage
arr = [1, 2, 3, 4]
mylst = LazySegmentTreeWithSetUpdate(arr)
mylst.update(0, 3, 0)
print("updating range from", 0, "to", 3, "with value:", 0)
mylst.update(2, 3, 2)
print("updating range from", 2, "to", 3, "with value:", 2)
mylst.update(0, 2, 10)
print("updating range from", 0, "to", 2, "with value:", 10)
ans = mylst.query(1, 2)
print("sum in range:", 1, "to", 2, ":", ans)
if __name__ == "__main__":
LazySegmentTreeWithSetUpdate.how_to_use()
# This code is contributed by arindam369