Unit 6: Questions About Ranges

A million numbers and a million questions

"What do positions 2 to 4 add up to?" — asked a million times, about a million numbers.

Adding them up each time is a million times a million. There is a way to answer every question with one subtraction.

Store the running total instead

prefix[i] is the total of everything before position i. Build it in one pass.

Then the total of positions a to b is prefix[b + 1] - prefix[a]: everything up to the end of the range, minus the part in front of it.

Python

numbers = [3, 1, 4, 1, 5, 9, 2, 6]

prefix = [0]
for x in numbers:
    prefix.append(prefix[-1] + x)

print(prefix)
print(prefix[5] - prefix[2])
print(numbers[2] + numbers[3] + numbers[4])

It prints

[0, 3, 4, 8, 9, 14, 23, 25, 31]
10
10

Why the extra 0 at the front

It is there so that a range starting at position 0 works without a special case. prefix[0] is "the total of nothing", which is 0.

Getting rid of special cases by adding a harmless extra slot is a habit that saves a lot of debugging under time pressure.

The same trick in two dimensions

For a grid, prefix[r][c] holds the total of the whole rectangle above and to the left. A rectangle sum is then four lookups: the big block, minus the strip above, minus the strip to the left, plus the corner that got taken away twice.

That last plus is where everybody makes the mistake.

Python

grid = [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9],
]

prefix = [[0] * 4 for r in range(4)]
for r in range(3):
    for c in range(3):
        prefix[r + 1][c + 1] = (grid[r][c] + prefix[r][c + 1]
                                + prefix[r + 1][c] - prefix[r][c])

def rectangle(r1, c1, r2, c2):
    return (prefix[r2 + 1][c2 + 1] - prefix[r1][c2 + 1]
            - prefix[r2 + 1][c1] + prefix[r1][c1])

print(rectangle(0, 0, 2, 2))
print(rectangle(1, 1, 2, 2))

It prints

45
28

One update ruins everything

Change position 3 and every prefix from there on is wrong. Rebuilding costs n, so a million updates cost a million million again.

What is needed is a structure where both jobs — change one number, total a range — cost about log n.

Boxes that hold overlapping chunks

A Fenwick tree is one list where box i holds the total of a chunk ending at i, and the chunk's length is the lowest set bit of i — that i & -i from unit 3.

So box 8 holds eight numbers, box 6 holds two, box 5 holds one. Any prefix is a handful of boxes, and any position sits in only a handful of boxes.

Python

for i in range(1, 9):
    print(i, bin(i), i & -i)

It prints

1 0b1 1
2 0b10 2
3 0b11 1
4 0b100 4
5 0b101 1
6 0b110 2
7 0b111 1
8 0b1000 8

Two loops, four lines

To add to a position, walk up the boxes that cover it: i += i & -i. To total a prefix, walk down: i -= i & -i.

Both loops knock out a bit each time, so both take about log n steps.

Python

n = 8
tree = [0] * (n + 1)

def add(position, value):
    i = position + 1
    while i <= n:
        tree[i] += value
        i += i & -i

def prefix(count):
    total = 0
    i = count
    while i > 0:
        total += tree[i]
        i -= i & -i
    return total

numbers = [3, 1, 4, 1, 5, 9, 2, 6]
for i in range(n):
    add(i, numbers[i])

print(prefix(8), prefix(5) - prefix(2))
add(3, 10)
print(prefix(8), prefix(5) - prefix(2))

It prints

31 10
41 20

Read it, use it, and do not re-derive it under pressure

Nobody works out the box layout from scratch in a contest. This is a machine you learn once and type from memory — four lines, and the +1 because box 0 would never move.

What you do have to know is when to reach for it: range totals with changing values.

Now the question is a minimum, and subtraction dies

Prefix sums work because a sum can be undone: take the total to here, take away the total to there. There is no such move for a smallest.

Knowing the smallest of the first 7 and the smallest of the first 2 tells you nothing about positions 2 to 6. This is not slow — it is impossible.

So build a pyramid instead

A segment tree stores the numbers along the bottom row, and every box above holds the smallest of the two under it. The top box is the smallest of everything.

Any range is covered by at most about 2 log n of those boxes, and changing one number only touches the boxes above it.

Python

numbers = [3, 1, 4, 1]
level = numbers
while len(level) >= 1:
    print(level)
    if len(level) == 1:
        break
    level = [min(level[i], level[i + 1]) for i in range(0, len(level), 2)]

It prints

[3, 1, 4, 1]
[1, 1]
[1]

Stored flat, walked from the bottom

The pyramid lives in one list of twice the size: position i of the data sits at size + i, and the box above any box i is i // 2.

To change a number, write it at the bottom and mend every box above it. To ask a range, walk both ends upwards, grabbing any box that sticks out.

Python

size = 8
big = 10 ** 9
tree = [big] * (2 * size)

def update(position, value):
    i = position + size
    tree[i] = value
    i //= 2
    while i >= 1:
        tree[i] = min(tree[2 * i], tree[2 * i + 1])
        i //= 2

def smallest(left, right):
    a = left + size
    b = right + size
    best = big
    while a < b:
        if a % 2 == 1:
            best = min(best, tree[a])
            a += 1
        if b % 2 == 1:
            b -= 1
            best = min(best, tree[b])
        a //= 2
        b //= 2
    return best

numbers = [3, 1, 4, 1, 5, 9, 2, 6]
for i in range(8):
    update(i, numbers[i])

print(smallest(0, 8), smallest(2, 5), smallest(5, 8))
update(3, 0)
print(smallest(2, 5))

It prints

1 1 2
0

The range is half open, and that is on purpose

smallest(2, 5) means positions 2, 3 and 4 — the right end is not included. Every walk-up-the-tree structure is written this way, because "the right end is where the next range starts" makes the arithmetic come out even.

Decide which convention you are using before you type, and then never mix them.

Three structures, one question each

Nothing changes, want a sum → prefix sums. Build once, answer in one subtraction, and never simpler than that.

Things change, want a sum → Fenwick tree.

Want a minimum, a maximum, or anything that cannot be undone → segment tree.

Counting the cost properly

With a million numbers and a million questions: prefix sums cost about two million steps in total. The naive loop costs about a million million.

A Fenwick tree costs about twenty per job, so forty million for a million updates and a million questions. All three of those numbers decide contests.

Python

import math

n = 1000000
q = 1000000
log = math.floor(math.log2(n)) + 1

print("naive", q * n)
print("prefix", n + q)
print("fenwick", (n + q) * log)

It prints

naive 1000000000000
prefix 2000000
fenwick 40000000

The trap: an update you did not notice

Half the mistakes here are reading the problem. A single sentence — "then the shop restocks item 4" — is the difference between prefix sums and a Fenwick tree.

So read the job list first. If any job changes a value, the prefix array is already dead.