Embercode
✕
Segment TreesLesson 5 of 15

55 min read

Module assignment: Scoreboard range service in Python

The project in one paragraph

A small league runs a scoreboard. Each player has a score, and the scores change during a match. The league website asks two kinds of question about a block of players. It asks for the total score of the block, and it asks for the lowest score in the block. The site also sends two kinds of change. It can set one player's score, and it can add a bonus to every player in a block. Your job is to write the Python program that answers these commands quickly. The program must handle a long list of players and a long list of commands, so a plain loop for each command will be too slow.

This assignment uses everything in the module. You will build a segment tree, store two values in each node, answer range queries, change single items, and apply a lazy range add. Work through it in the order given below. Each step builds on the one before it, and the earlier steps give you partial credit even if you do not finish the later ones.

What you must build

Write one Python program in a file called solution.py. The program reads the input from standard input, which means the text you pipe into it or type into a terminal, and prints the answers to standard output. The first version of your program may be a plain loop, which is fine for a first check. The final version must use a segment tree with lazy propagation, so that each command takes O(log n) time.

The scoreboard has n players, numbered from 0 to n - 1. Each player has an integer score. The program stores the scores in a tree, and every command reads or changes a block of players. The block is always given as a left index l and a right index r, and both ends are included. The answers are sums or minimums of that block, so the tree must store both of those for each node.

Input and output format

The input has four parts, in this order.

Line 1 has one number, n, the number of players. It is between 1 and 200000.

Line 2 has n scores, separated by spaces. Each score is a whole number from -1000000000 to 1000000000.

Line 3 has one number, q, the number of commands. It is between 1 and 200000.

Each of the next q lines is one command. There are four kinds.

  • SUM l r asks for the total of the scores from l to r. The program prints one line with the total.
  • MIN l r asks for the lowest score from l to r. The program prints one line with the minimum.
  • SET i v changes the score of player i to v. The program prints nothing.
  • ADD l r v adds v to the score of every player from l to r. The value v can be negative. The program prints nothing.

The program prints one line for each SUM and each MIN, in the same order as the commands. Totals can be large, but Python integers grow as needed, so you do not need to think about overflow. The task is still only correct if every value is read as a whole number, not as text.

Examples

The first example is the one the league used to test its site. Its output has four lines, one for each SUM and MIN command.

6
4 9 2 7 5 8
6
SUM 0 5
MIN 1 3
SET 2 1
MIN 0 5
ADD 2 4 10
SUM 2 4

The output is:

35
2
1
43

Check each line by hand. The first SUM 0 5 adds all six scores, which is 4 + 9 + 2 + 7 + 5 + 8 = 35. The MIN 1 3 looks at 9, 2 and 7, and the lowest is 2. The SET 2 1 changes player 2 from 2 to 1, so the scores are now 4 9 1 7 5 8. The MIN 0 5 then finds 1. The ADD 2 4 10 adds 10 to players 2, 3 and 4, which makes the scores 4 9 11 17 15 8. The last SUM 2 4 adds 11 + 17 + 15 = 43.

The second example uses negative scores. Its output has three lines.

3
5 5 5
5
ADD 1 2 -7
SUM 0 2
MIN 0 2
SET 0 10
MIN 0 2

The output is:

1
-2
-2

After ADD 1 2 -7 the scores are 5 -2 -2. The sum is 5 - 2 - 2 = 1, and the minimum is -2. The SET 0 10 changes the first score to 10, but the minimum is still -2 because the other two scores are lower than 10.

Build your own small tests as well. A one-player list is a good test, because it checks the leaf case. A command with l equal to r is another good test. A large test, with 200000 commands, checks the speed.

To run a program on an input file, save the input as input.txt and type python3 solution.py < input.txt. The < sign sends the file's text to the program as if you had typed it.

Requirements checklist

Tick these off as you finish each one. A program that meets every item in the first six groups is a complete solution.

  • The program reads n, the n scores, q, and the q commands in that order.
  • Every number is read with int(), so the comparison and the arithmetic work on whole numbers.
  • The tree stores the sum of each node in one list and the minimum of each node in another.
  • SUM l r prints the sum of the block from l to r.
  • MIN l r prints the minimum of the block from l to r.
  • SET i v changes one score and recomputes every node on the path to that player.
  • ADD l r v adds v to a whole block with a lazy note, so the update costs O(log n).
  • Each query pushes the pending note down before it reads a child, and the update does too.
  • The program prints the answers in the same order as the commands, with one answer per line.
  • A one-player list works, and so does a command where l equals r.
  • A test with 200000 commands finishes in a time you have measured and written down. Python is slower than Java here, so the time is not the one a Java program would reach.
  • The code is readable. Each method has one clear job, and the names say what they hold.

Grading rubric

The grade depends on how many commands the program answers correctly and how fast it does so. Your work is measured against three levels.

A 60% solution answers the commands correctly for small inputs. It may use a plain loop for each command. It handles SUM, MIN and SET, and it may not handle ADD at all. Its output is correct on the two examples above. It may be slow on large inputs. A 60% solution proves that you understand the commands and the input format.

An 80% solution uses a segment tree for SUM, MIN and SET. Each of those commands costs O(log n). It answers every command correctly on medium inputs, including the negative numbers. It may still use a loop for ADD, which is correct but slow. An 80% solution proves that you can build and query the tree and handle point updates.

A 100% solution uses a lazy segment tree for all four commands. The tree stores both the sum and the minimum in each node. The ADD command uses a pending note, and every query and update pushes notes before it reads a child. The program is correct on every test, including a test with 200000 commands. A 100% solution proves that you can combine the module's ideas into one working system.

Graders also look at the code itself. A correct program that is hard to read loses some marks, and a program with an unused variable or a copied block of code loses a few more. Write small methods, name them after their job, and keep the lazy helpers together.

Plan the work in steps

Do the work in the order below. Each step has a clear test, so you know when it is right before you move on.

Step one is reading the input. Read n, the scores, q, and then each command. Print each command's name and numbers. Check that the list of commands matches the file you typed. Do not build a tree yet.

Step two is the plain version. Write each command as a loop over the list. SUM adds the block with sum, MIN finds the smallest with min, SET changes one item, and ADD adds to each item in the block. Check the output against the first example. This version is slow, but it is the reference for every later step.

Step three is the sum tree with point updates. Build a tree that stores sums, and answer SUM queries with it. Make SET a point update that recomputes the path. Run the first example, and then your own test with a few commands. Compare with the plain version.

Step four is the minimum tree. Add a second list for the minimum of each node. Answer MIN queries with it, and make SET update the minimum on the path too. The no-overlap value for MIN must be float("inf"), because a missing block must never win a minimum.

Step five is the lazy range add. Add the note list and the apply and push helpers. Make ADD a range update. Apply the note to the sum and the minimum, because adding v to every item also adds v to the lowest item of the block. Push the note before any recursion into a child, in both SUM and MIN queries and in the update.

Step six is the final check. Compare your program with the plain version on random commands. Use a small size, such as 20 players and 200 commands, and print the first command where the two programs disagree. Then time the program on the largest input.

Hints

Open a hint only when you are stuck, and try the step once more before you read the next one.

Hint 1: the tree has two lists

Use self.sum and self.mn as two lists, both of size 4 * n. Build them with the same recursive shape as the lessons. For a leaf, set both to the item's value. For a parent, set self.sum to the sum of the children and self.mn to the smaller of the two children's minimums. Write one _pull method that recomputes both values from the children, and call it after every change.

Hint 2: the apply helper handles both values

Write _apply_add(node, l, r, v) that does three things. It adds v * (r - l + 1) to the sum of the node, adds v to the minimum of the node, and adds v to the note of the node. The minimum grows by v because every item in the block grows by v, and the smallest item grows too. Use this helper both in the ADD update and in the push.

Hint 3: push before you go below a partly covered node

A _push method sends the note to both children with _apply_add, and then clears the note. Call it in _add_range before the recursion, and call it in _query_sum and _query_min before their recursion. Do not call it on a node that is fully inside the range or fully outside, because those stop the search.

Hint 4: SET changes one leaf

For SET i v, walk down to the leaf for i. Push on the way down, so the path is correct. Set the leaf's sum and minimum to v, and clear its note. Then pull each parent on the way back up. A point set replaces a value rather than adding to it, so write it as its own method and do not reuse the add helper.

Hint 5: a missing block for each query

For _query_sum, a block with no overlap returns 0. For _query_min, it returns float("inf"). If you use 0 for _query_min, any query that does not cover the whole list will return 0 and the answers will be wrong. Check this with the first example, where MIN 1 3 covers only part of the list.

Hint 6: read the input the same way every time

Read the whole input at once with sys.stdin.read().split(), which gives a list of words. Walk through the list with a position counter. Compare the command word with ==, and convert each number with int(). The first example is a good test of the reading code on its own.

Reference solution

Try to build the program yourself before you open the solution. The reference solution below is a complete, correct program that handles all four commands with a lazy segment tree. It reads the input with sys.stdin.read(), collects the answers in a list and prints them once at the end.

Reference solution

import sys
 
 
class ScoreBoard:
    def __init__(self, a):
        self.n = len(a)
        size = 4 * self.n
        self.sum = [0] * size
        self.mn = [0] * size
        self.lazy = [0] * size
        self._build(a, 1, 0, self.n - 1)
 
    def _build(self, a, node, l, r):
        if l == r:
            self.sum[node] = a[l]
            self.mn[node] = a[l]
            return
        mid = (l + r) // 2
        self._build(a, 2 * node, l, mid)
        self._build(a, 2 * node + 1, mid + 1, r)
        self._pull(node)
 
    # Recompute a parent from its two children.
    def _pull(self, node):
        self.sum[node] = self.sum[2 * node] + self.sum[2 * node + 1]
        self.mn[node] = min(self.mn[2 * node], self.mn[2 * node + 1])
 
    # Add v to every item of a fully covered segment. Store the note in lazy.
    def _apply_add(self, node, l, r, v):
        self.sum[node] += v * (r - l + 1)
        self.mn[node] += v
        self.lazy[node] += v
 
    def _push(self, node, l, r):
        if self.lazy[node] != 0:
            mid = (l + r) // 2
            self._apply_add(2 * node, l, mid, self.lazy[node])
            self._apply_add(2 * node + 1, mid + 1, r, self.lazy[node])
            self.lazy[node] = 0
 
    def add_range(self, ql, qr, v):
        self._add_range(1, 0, self.n - 1, ql, qr, v)
 
    def _add_range(self, node, l, r, ql, qr, v):
        if qr < l or r < ql:
            return
        if ql <= l and r <= qr:
            self._apply_add(node, l, r, v)
            return
        self._push(node, l, r)
        mid = (l + r) // 2
        self._add_range(2 * node, l, mid, ql, qr, v)
        self._add_range(2 * node + 1, mid + 1, r, ql, qr, v)
        self._pull(node)
 
    def set_point(self, i, v):
        self._set_point(1, 0, self.n - 1, i, v)
 
    def _set_point(self, node, l, r, i, v):
        if l == r:
            self.sum[node] = v
            self.mn[node] = v
            self.lazy[node] = 0
            return
        self._push(node, l, r)
        mid = (l + r) // 2
        if i <= mid:
            self._set_point(2 * node, l, mid, i, v)
        else:
            self._set_point(2 * node + 1, mid + 1, r, i, v)
        self._pull(node)
 
    def query_sum(self, ql, qr):
        return self._query_sum(1, 0, self.n - 1, ql, qr)
 
    def _query_sum(self, node, l, r, ql, qr):
        if qr < l or r < ql:
            return 0
        if ql <= l and r <= qr:
            return self.sum[node]
        self._push(node, l, r)
        mid = (l + r) // 2
        return self._query_sum(2 * node, l, mid, ql, qr) + self._query_sum(2 * node + 1, mid + 1, r, ql, qr)
 
    def query_min(self, ql, qr):
        return self._query_min(1, 0, self.n - 1, ql, qr)
 
    def _query_min(self, node, l, r, ql, qr):
        if qr < l or r < ql:
            return float("inf")
        if ql <= l and r <= qr:
            return self.mn[node]
        self._push(node, l, r)
        mid = (l + r) // 2
        return min(self._query_min(2 * node, l, mid, ql, qr), self._query_min(2 * node + 1, mid + 1, r, ql, qr))
 
 
def main():
    data = sys.stdin.read().split()
    pos = 0
    n = int(data[pos])
    pos += 1
    a = [int(x) for x in data[pos:pos + n]]
    pos += n
    board = ScoreBoard(a)
 
    q = int(data[pos])
    pos += 1
    out = []
    for _ in range(q):
        op = data[pos]
        pos += 1
        if op == "SUM":
            out.append(str(board.query_sum(int(data[pos]), int(data[pos + 1]))))
            pos += 2
        elif op == "MIN":
            out.append(str(board.query_min(int(data[pos]), int(data[pos + 1]))))
            pos += 2
        elif op == "SET":
            board.set_point(int(data[pos]), int(data[pos + 1]))
            pos += 2
        else:
            board.add_range(int(data[pos]), int(data[pos + 1]), int(data[pos + 2]))
            pos += 3
    if out:
        print("\n".join(out))
 
 
if __name__ == "__main__":
    main()

Why the reference is built this way

The reference follows the rules from the lessons, and it is worth reading each part against them. The build fills the sums and the minimums from the children, as in the first two lessons. The _pull method is the one place that recomputes a parent, so every change calls it. The _apply_add method is the only place that changes a node's values by a note, so the sum uses the length and the minimum uses the plain value.

The _set_point method pushes on the way down, which keeps the path correct before it changes the leaf. It clears the leaf's note, because a leaf has no children to deliver it to. The update and query methods push only when they recurse into a partly covered node. That is the rule from the lazy lesson, and it is why the reference answers stay correct.

The reading code converts each value with int(), and it compares the command word with ==. The answers go into a list and are printed once at the end. Printing once is faster than printing each line, and the time is important for the large test.

Speed is the one place where Python differs from a compiled language. Each recursive call is a function call with a real cost, so a tree with 200000 items and 200000 commands makes millions of calls. On the machine used to check this lesson, a random input with 200000 players and 200000 commands took about 4.8 seconds. Your time will differ with the machine and the Python version. The point of the rubric is the method, and a program that is correct but slow still earns the correctness marks.

Extension: a stretch goal

If you finish all the checklist items, try one extension. Add a command CNT l r t that counts how many scores in the block from l to r are at least t. This needs a different kind of node. One option stores a sorted list of values in every node, which is called a merge sort tree. Python's bisect module does the binary search inside each sorted list. Think about the cost before you write any code. The sorted lists use O(n log n) memory in total, and a query binary searches one list per visited node, which makes it about O(log^2 n). Explain the trade-off in a short paragraph, even if you do not write the code.

Submitting and self-check

Before you submit, run the program on both examples and check the output line by line. Then run it on three small random tests that you made yourself, and compare each with the plain version from step two. Finally, run the largest test you can make and check that it finishes in a reasonable time on your computer. Write the time down.

Write one short paragraph that says which rubric level you reached and why. Name the one command you are least sure about. Explain in plain words what the _pull and _push methods do, as if you were teaching a classmate. That paragraph is part of the grade, because a program that works but cannot be explained is not yet understood.

Module summary

This module covered the segment tree from first principles, through the following steps.

  • Range queries with changing data need a structure better than a loop or prefix sums.
  • A segment tree stores one answer per segment, builds in O(n), and answers queries and point updates in O(log n).
  • A query uses three cases: no overlap, full overlap and partial overlap, and only partial nodes recurse.
  • The combine rule must be associative and have a neutral value. Sum, minimum, maximum and gcd all fit.
  • Lazy propagation applies a range update with a note, pushes the note down before reading a child, and keeps the total cost at O(log n).
  • A Fenwick tree is a compact alternative for sums with point updates, while a segment tree is the flexible choice.
Practice this topic: coding problems for Segment Trees →

Stuck? Ask the tutor

Ask about anything in this lesson. Answers are short and use simple words.

Sign in to ask the tutor.