Range assign and range add with lazy tags in Python

Why this matters

Imagine a row of 8 seats, each with a price. A shop can do two kinds of bulk change. It can set every seat in a block to the same price, which is an assignment. It can raise every seat in a block by the same amount, which is an addition. After each change, the shop may ask for the total price of any block.

If the shop changes one seat at a time, a block of 1000 seats costs 1000 steps per change. A segment tree is a tree that can change or total a whole block in a few steps. Each node of the tree covers a range of positions. The root covers all of them, and the two children cover the left half and the right half. The leaves, the nodes with no children, cover single positions.

The trick that makes the tree fast is a lazy tag. A tag is a note that says "every value in this node's range must change in this way, and the children have not been told yet." The node's own total is fixed right away. The children are updated only when someone asks about or changes a part of the range that lies below the node.

The earlier lesson on lazy propagation introduced this idea for addition alone. This lesson adds assignment as well. Once there are two kinds of change, the order of the changes matters, and a tag that stores them carelessly gives wrong answers without an error message. Most of the lesson is about keeping that order right.

You should be comfortable with recursion, classes, and the dataclass decorator. A dataclass is a short way to write a class whose main job is to hold a few values.

The idea in a small example

Take the list 1 2 3 4. Add 10 to every value. Then ask for the sum of all four values.

With a plain list, the change touches four values, and the sum is 11 + 12 + 13 + 14 = 50.

With a segment tree, the root covers the whole list. The root's sum is changed at once. The root's sum was 10 and becomes 10 + 4 * 10 = 50. The root's tag says "add 10 to each of your values". The two children, which cover 1 2 and 3 4, still show their old sums and do not have the new values yet.

Now ask for the sum of the whole list. The root covers the full range of the question, so the tree answers from the root alone, 50, without looking at the children.

Ask for the sum of just the first position. The root does not fit inside that question, so the root must hand its tag down. It gives the tag to both children. The left child's sum changes from 3 to 23. The left child's own tag is now "add 10 to each of your values". The tree then goes down one more level and gives the leaf for position 0 the change, so that leaf's value becomes 11. The answer is 11.

The tag was carried down only along the path to the answer. The other branches never had to be touched.

One tag, two kinds of change

A tag must describe both kinds of change, so it has three parts. Here is the rule for what a tag does to each value x in its range:

• If the tag has no assignment, the value becomes x + add.
• If the tag has an assignment value, the value becomes value + add. The old value x is ignored.

The dataclass below stores these three parts.

from dataclasses import dataclass
 
 
@dataclass(frozen=True)
class Tag:
    set: bool
    value: int
    add: int
 
 
NONE = Tag(False, 0, 0)

The frozen=True option makes each tag read-only. A tag is then a value you can compare and share, and no code can change one in place by mistake. NONE is a tag that changes nothing. It is the starting state of every node.

Here are two examples. The tag Tag(False, 0, 10) means "add 10". The tag Tag(True, 5, 0) means "make every value 5". The tag Tag(True, 5, 2) means "make every value 5, then add 2", so every value becomes 7.

When a tag is applied to a node that covers length positions, the node's sum changes like this:

• If the tag has an assignment, the new sum is (value + add) * length. The old sum is thrown away, because every value is replaced.
• If the tag has no assignment, the new sum is the old sum plus add * length.

In code, the method looks like this. It is part of the full program later in the lesson.

def _apply_to_node(self, node, l, r, tag):
    length = r - l + 1
    if tag.set:
        self.sum[node] = (tag.value + tag.add) * length
    else:
        self.sum[node] += tag.add * length
    self.lazy[node] = self.lazy[node].then(tag)

The last line is the important one. It combines the node's existing tag with the new tag. That combining step is the subject of the next section.

Combining two tags: the order rule

A node may receive several changes before anyone asks about its children. Say the node already has a pending tag f, which is an older change. Then a new change g arrives. The node must now hold one tag that does what f does and then what g does. Call the combined tag "f then g".

There are three cases.

Case 1: the new change is an assignment. The assignment overwrites every value, so whatever f did is gone. The combined tag is just g.

Example: the old tag is "add 9", and the new change is "assign 5". Start with the values 1 and 2. After the add they are 10 and 11. After the assignment both are 5. The combined tag is Tag(True, 5, 0), and the add 9 has disappeared.

Case 2: the new change is an addition, and the old tag has an assignment. The assignment sets the value, and the addition is applied after it. So the combined tag keeps the assignment and adds the new amount to its add part.

Example: the old tag is "assign 5", which is Tag(True, 5, 0). The new change is "add 2". Every value is 5 and then becomes 7. The combined tag is Tag(True, 5, 2).

Case 3: the new change is an addition, and the old tag has no assignment. Two additions in a row add up, so the add parts are summed.

Example: the old tag is "add 3" and the new change is "add 4". The combined tag is "add 7".

In Python, the rule is one method on the dataclass:

def then(self, later):
    if later.set:
        return later
    return Tag(self.set, self.value, self.add + later.add)

If later has an assignment, the old tag does not matter, so later is returned. If later has no assignment, its addition is added to the old tag's add part, and the old tag's assignment, if any, is kept. Because the tag is frozen, then returns a new Tag instead of changing the old one.

Why one combined tag keeps the order right

The program stores one combined tag per node, not separate lists of assignments and additions. This is what keeps the order correct. Every time a change arrives, then puts it after all the changes already stored. The stored tag always means "the earlier changes, then the later ones".

There is one more rule about order, and it is the one that most bugs break. A node's tag is always newer than the tags of its children. The reason is that a change reaches a child only after the parent's tag has been pushed down. When a change is pushed from a node to its children, the child combines its own tag with the node's tag using then. The node's tag is the later change, so it goes second.

Here is the same pattern written in the push method:

def _push(self, node, l, r):
    if self.lazy[node] == NONE:
        return
    mid = (l + r) // 2
    self._apply_to_node(2 * node, l, mid, self.lazy[node])
    self._apply_to_node(2 * node + 1, mid + 1, r, self.lazy[node])
    self.lazy[node] = NONE

The _apply_to_node call combines the child's tag with the node's tag, so the child receives the node's changes after its own. After both children have the tag, the node's tag is cleared, because the children now carry it.

The node's tag Tag(True, 5, 2) goes down to a child that holds the older tag "add 9". The child's combined tag becomes Tag(True, 5, 2), because the assignment in the node's tag overwrites the child's add. The child's values become 7, which is what they would be if the changes had been applied in the order they happened.

The full program

The program below builds a segment tree from a list of values. It then runs a list of operations. Each operation is one of three kinds:

• ("assign", l, r, v) sets every value from position l to position r to v.
• ("add", l, r, d) adds d to every value from position l to position r.
• ("sum", l, r) prints the total of the values from position l to position r.

The positions start at 0 and both ends are included. The operations in this program are the first worked example. To try another example, change the list of values and the list of operations.

The output is:

sum 0 4 = 15
sum 0 4 = 46
sum 2 2 = 12

The _update and _query methods follow the same pattern as any segment tree. If a node's range is fully inside the question, the method works on the node and stops. If the range is outside the question, the method stops and returns nothing. If the range is partly inside, the method pushes the node's tag down, then recurses into both children. The push is what makes the children correct before the method looks at them.

The _update method finishes by recomputing the node's sum from its children. This is needed because the children may have changed in part.

The values list and the operations list are the only things to change for a new test. The rest of the program stays the same.

Check the second line by hand. After the assignment the values are 1 10 10 10 5. After adding 2 to each, they are 3 12 12 12 7. The sum is 3 + 12 + 12 + 12 + 7 = 46. The last line asks for position 2 alone, which is 12.

Worked example two: assign, then add

This example shows case 2 from the order rule. The list is 1 2 3 4. The program assigns 5 to all positions, adds 2 to all positions, and sums. Then it adds 4 to positions 0 and 1, assigns 1 to position 0, and sums again.

For this example, set values = [1, 2, 3, 4] and operations to these six steps:

operations = [
    ("assign", 0, 3, 5),
    ("add", 0, 3, 2),
    ("sum", 0, 3),
    ("add", 0, 1, 4),
    ("assign", 0, 0, 1),
    ("sum", 0, 3),
]

The output is:

sum 0 3 = 28
sum 0 3 = 26

After the first two changes every value is 7, so the sum is 28. After the addition of 4 to positions 0 and 1, the values are 11 11 7 7. Then position 0 is assigned 1, so the values are 1 11 7 7, and the sum is 26.

Worked example three: add, then assign

This example shows case 1. The list is 1 2 3 4. The program adds 9 to all positions, then assigns 5 to all of them. Set operations to [("add", 0, 3, 9), ("assign", 0, 3, 5), ("sum", 0, 3)].

The output is sum 0 3 = 20. The add of 9 is overwritten by the assignment, so each value is 5 and the sum is 20. If the add were kept on top of the assignment, every value would be 14 and the sum would be 56.

Worked example four: a tag that is split between children

The last example uses a range that covers only part of a node, so the tag has to go down to the children. The list is 1 2 3 4. The program assigns 5 to positions 0 and 1, then adds 1 to positions 1 and 2, and sums everything. Set operations to [("assign", 0, 1, 5), ("add", 1, 2, 1), ("sum", 0, 3)].

The output is sum 0 3 = 19.

The by-hand check goes like this. After the assignment the values are 5 5 3 4. After the addition on positions 1 and 2, the values are 5 6 4 4. The sum is 5 + 6 + 4 + 4 = 19.

The second change covers position 1 of the left half and position 2 of the right half. In the tree, the left half 0 1 has an assignment waiting in its tag. Before the add can land on position 1, that assignment must be pushed down to the two leaves of the left half. The push step does this, and the add then changes the leaves for positions 1 and 2 only.

Checking the tree against a plain list

A segment tree can be wrong in a way that a few hand examples miss. The safe check is to run many random operations on both the tree and a plain list, and compare every sum. The program below does this for 500 random lists, each with up to 30 operations. It uses a fixed random seed, so its output is the same on every run.

The output is:

random trials: 500, queries checked: 2491, mismatches: 0

The check covers 2491 sums across the random operations. Every sum matched the plain list.

Common mistakes

Forgetting to push before reading the children

The _query method must push a node's tag before it reads the children. If the push line is removed, the query reads children whose values have not been changed yet.

Take the program above and delete the line self._push(node, l, r) inside _query. Then run worked example one again. The output becomes:

sum 0 4 = 15
sum 0 4 = 46
sum 2 2 = 10

The answer for position 2 is 10, when the correct value is 12. Python prints this wrong number with no error. The first two sums are correct only because they are answered from nodes that fully cover the question.

Clearing the tag after the push

The _push method must set the node's tag back to NONE after passing it down. Otherwise the same tag is passed down again the next time the node is pushed, and the children change twice.

Take the program above and delete the line self.lazy[node] = NONE inside _push. Then run this input: a list of [1, 2], then add 0 1 3, sum 0 0, add 1 1 0, sum 0 0. The correct output is sum 0 0 = 4 twice. Position 0 starts at 1, gets 3 added, and stays at 4. Without the clear, the second push adds the 3 again, and the output is:

sum 0 0 = 4
sum 0 0 = 10

The second answer is 10, which is wrong. The mistake only shows up when the same node is pushed more than once, so small tests often miss it.

Pushing the two parts of a tag in the wrong order

A common design keeps two separate lists, one for assignments and one for additions. That design is easy to get wrong. If the push passes the addition down first and then the assignment, the children are given the change in the wrong order. The assignment must go down first, and the addition after it.

The program below stores the two parts separately and pushes them in the wrong order. Its output shows the problem.

The output is:

whole range sum: 14
value at index 0: 5
value at index 1: 5

Follow the steps. The list is 1 2. The program assigns 5 to both positions, then adds 2 to both. Both values should now be 7, and the sum should be 14. The root's sum is 14, so the whole-range question is right.

The question for index 0 forces the root to push. The wrong push adds 2 to the children first. The children become 3 and 4. Then it assigns 5, which overwrites them, and both children become 5. The correct answer for index 0 is 7, and the program prints 5.

The root's sum still says 14, while its children add up to 10. The tree now disagrees with itself. The fix is to push the assignment first and the addition second. A single combined tag, as in the main program, does this automatically, because the tag is always one value.

Asking for positions outside the list

The tree does not check that the question is inside the list. A query such as sum 2 9 on a list of length 4 does not fail. It answers for the part that exists.

Run the program with values = [1, 2, 3, 4] and the single operation ("sum", 2, 9). The output is sum 2 9 = 7. Positions 2 and 3 hold 3 and 4, so the answer is 7. There is no error, so the caller has to check the range before asking. The same is true for assign and add, which change only the part that exists.

Practice

1. A node has pending tag "add 3", then receives "assign 4", then receives "add 1". What single tag should the node hold? What does each value in the range become?

Show the answer

The single tag is Tag(True, 4, 1). The assignment of 4 overwrites everything before it, and the add of 1 comes after it, so each value becomes 4 + 1 = 5.

2. The list is 2 3 4 5. Assign 1 to positions 0 to 2, then add 10 to positions 1 to 3, then ask for the sum of all four positions. What is the answer?

Show the answer

The answer is 38. After the assignment the values are 1 1 1 5. After adding 10 to positions 1 to 3 they are 1 11 11 15, and the sum is 1 + 11 + 11 + 15 = 38.

3. A parent node has the tag "assign 2" pending. Its left child has the tag "add 5" pending from an earlier change. The parent is pushed down. What tag does the left child hold afterwards?

Show the answer

The left child holds Tag(True, 2, 0). The parent's assignment is newer than the child's add, so it overwrites the child's pending add. The child's values will be 2, and the old add of 5 disappears. Case 1 of the order rule gives this result.

4. The list is 1 2 3. Add 4 to positions 0 to 1, then assign 0 to positions 1 to 2, then ask for the sum of positions 0 to 2. What is the answer, and which part of the order rule decides it?

Show the answer

The answer is 5. After adding 4 to positions 0 and 1, the values are 5 6 3. The assignment of 0 to positions 1 and 2 gives 5 0 0, and the sum is 5. Case 1 decides it, because the assignment overwrites the earlier add on the positions it covers, and the add still applies to position 0, which the assignment did not cover.

5. Explain in two or three sentences why a node's tag must always be newer than the tags of its children.

Show the answer

A child's pending tag was stored before the parent's tag, because a parent is always pushed down before it receives a new change. If the parent's tag were older, applying it to the child would undo the newer change stored in the child. The rule keeps the changes in the order they happened.

Summary

• A lazy tag stores a pending change for a node's range, with a flag for assignment, an assignment value, and an addition amount.
• When a new change is combined with an old tag, an assignment replaces the old tag, and an addition adds to the old tag's amount.
• Push a node's tag to its children before reading or changing them, then clear the node's tag.
• A node's tag is always newer than its children's tags, so pushing it down keeps the changes in the order they happened.