Lazy propagation

Adding a value to every position of a range would change up to n leaves of a segment tree. Lazy propagation changes only the O(log n) nodes that cover the range, and leaves a note on each saying what the nodes below it still owe. The note is passed down only when someone needs to look below.

Changing a whole range

Now the operations are add(l, r, v), add v to every value of a[l:r], and total(l, r), the sum of a[l:r], mixed in any order. The difference list of lesson 1 handles the adds, but only when every question comes after every add. A segment tree with single-value changes would need r − l changes per add.

The idea: a range add visits the same O(log n) nodes as a range question. A node that lies completely inside the range knows exactly how its own sum changes: by v times the number of positions it covers. Its children change too, but nobody has asked about them yet.

Tags

So give every node a second number, a tag (also called a lazy value): an amount that must still be added to every position below the node, but has not reached its children yet.

A range add then works like a range question. When a node lies completely inside the range, add v times its length to its sum, add v to its tag, and stop there. When it is only partly inside, go to both children and, on the way back, recompute its sum from theirs. Step through “Add 3 to a[1:7]”: two nodes get tags, and the cells below them keep their old values.

Pushing a tag down

Before any operation goes below a node with a tag, it pushes the tag down: each child adds the tag times its own length to its sum, and adds the tag to its own tag. Then the node’s tag is cleared. Now the children are correct, and the walk can go on. Step through “Then sum a[5:7]”.

A push costs O(1) and happens only at a node that the walk visits anyway. So the walk still visits O(log n) nodes, and both operations stay O(log n), however large the ranges. A tag can wait below a node for a long time, and several adds can pile up in one tag, because adding is the same whatever the order: +2 then +3 is +5.

The code

class RangeTree:
    def __init__(self, n):
        self.n = n
        self.sum = [0] * (4 * n)
        self.tag = [0] * (4 * n)

    def _apply(self, k, length, v):    # node k gets +v everywhere
        self.sum[k] += v * length
        self.tag[k] += v

    def _push(self, k, lo, mid, hi):
        if self.tag[k]:
            self._apply(2 * k, mid - lo, self.tag[k])
            self._apply(2 * k + 1, hi - mid, self.tag[k])
            self.tag[k] = 0

    def add(self, l, r, v, k=1, lo=0, hi=None):
        if hi is None:
            hi = self.n                  # node k covers a[lo:hi]
        if r <= lo or hi <= l:
            return                       # no overlap
        if l <= lo and hi <= r:
            self._apply(k, hi - lo, v)   # completely inside: tag it
            return
        mid = (lo + hi) // 2
        self._push(k, lo, mid, hi)
        self.add(l, r, v, 2 * k, lo, mid)
        self.add(l, r, v, 2 * k + 1, mid, hi)
        self.sum[k] = self.sum[2 * k] + self.sum[2 * k + 1]

    def total(self, l, r, k=1, lo=0, hi=None):
        if hi is None:
            hi = self.n
        if r <= lo or hi <= l:
            return 0
        if l <= lo and hi <= r:
            return self.sum[k]
        mid = (lo + hi) // 2
        self._push(k, lo, mid, hi)
        return (self.total(l, r, 2 * k, lo, mid)
                + self.total(l, r, 2 * k + 1, mid, hi))

Here the tree is stored top-down: node 1 covers a[0:n], and node k covering a[lo:hi] has children 2k and 2k + 1 covering the two halves. When n is not a power of two the numbering has gaps, and 4n places are always enough. The recursion is only about log₂ n deep, so Python’s recursion limit is no problem. Practice RNGQ-P5 asks for this class.

Other pairs of change and question work the same way, as long as a tag can be applied to a whole node in O(1) and two tags can be joined into one: “set every value to v” with sums (a later set replaces an earlier tag), or “add v” with minimum (the minimum grows by v, whatever the length).