added range update code

src/data_structures/sqrt_decomposition.md

@@ -110,6 +110,55 @@ For example, let's say we can do two types of operations on an array: add a give
 
 Finally, those two classes of problems can be combined if the task requires doing **both** element updates on an interval and queries on an interval. Both operations can be done with $O(\sqrt{n})$ complexity. This will require two block arrays $b$ and $c$: one to keep track of element updates and another to keep track of answers to the query.
 
+```py
+class SqrtDecomp:
+    def __init__(self, nums):
+        self.nums = nums
+        self.width = int(len(nums) ** 0.5) + 1
+        self.bn = (len(nums) // self.width) + 1  # number of blocks
+        self.blocks = [0] * self.bn  # precomputed sums
+        self.lazy = [0] * self.bn  # an additional lazy[block] is added when querying individual elements in a block
+        for i in range(len(nums)):
+            self.blocks[i // self.width] += nums[i]  # add to corresponding block
+
+    def update(self, i, j, diff):
+        first = (i // self.width) + 1  # first fully covered block
+        last = (j // self.width) - 1  # last fully covered block
+        if first > last:  # doesn't cover any blocks
+            for v in range(i, j + 1):
+                self.nums[v] += diff
+                self.blocks[v // self.width] += diff
+            return
+
+        for b in range(first, last + 1):  # blocks in the middle
+            self.lazy[b] += diff
+            self.blocks[b] += diff * self.width
+        for v in range(i, first * self.width):  # individual cells
+            self.nums[v] += diff
+            self.blocks[first - 1] += diff
+        for v in range((last + 1) * self.width, j + 1):
+            self.nums[v] += diff
+            self.blocks[last + 1] += diff
+
+    def query(self, i, j):
+        first = (i // self.width) + 1  # first fully covered block
+        last = (j // self.width) - 1  # last fully covered block
+        res = 0
+
+        if first > last:  # doesn't cover any blocks
+            for v in range(i, j + 1):
+                res += self.nums[v] + self.lazy[v // self.width]
+            return res
+
+        for b in range(first, last + 1):  # add value from blocks
+            res += self.blocks[b]
+        for v in range(i, first * self.width):  # add value from individual cells
+            res += self.nums[v] + self.lazy[first - 1]
+        for v in range((last + 1) * self.width, j + 1):
+            res += self.nums[v] + self.lazy[last + 1]
+        return res
+```
+
 There exist other problems which can be solved using sqrt decomposition, for example, a problem about maintaining a set of numbers which would allow adding/deleting numbers, checking whether a number belongs to the set and finding $k$-th largest number. To solve it one has to store numbers in increasing order, split into several blocks with $\sqrt{n}$ numbers in each. Every time a number is added/deleted, the blocks have to be rebalanced by moving numbers between beginnings and ends of adjacent blocks.
 
 ## Mo's algorithm