Contest Python: USACO Bronze

unit

Unit 8 — Prefix Sums

A contest scoreboard stores a huge row of point values, then asks thousands of questions about totals between two positions. Adding the same values again for every question is too slow. This unit builds cumulative totals once, answers each 1D range in constant time, then extends the same idea to rectangles of a grid.

Each idea comes as a short ladder — the core idea on a tiny array you can trace by hand, then one step up, then the full program that reads the real input from stdin.

You will learn

  • Build a 1D prefix-sum list in one pass.
  • Answer any range-sum question with a single subtraction.
  • Use the one-extra-entry convention to avoid off-by-one mistakes.
  • Build a 2D prefix grid and sum any rectangle with four corners.

Lesson 1 — 1D Prefix Sums

For an array a of length n, build a prefix array pre of length n + 1. The extra first entry pre[0] = 0 means "zero values have total zero", and pre[i] = pre[i - 1] + a[i - 1] is the sum of the first i values.

a = [4, 1, 3, 2]
pre = []
pre.append(0)
i = 0
while i < len(a):
    pre.append(pre[i] + a[i])
    i = i + 1
print(pre)

Output

[0, 4, 5, 8, 10]

Notice

pre carries one extra entry — pre[0] = 0, then each pre[i] is the running total of the first i values (here [0, 4, 5, 8, 10]).

left = 1
right = 3
print(pre[right + 1] - pre[left])

Output

6

Notice

The inclusive total of a[left..right] is pre[right + 1] - pre[left] — the + 1 on right keeps a[right] in the total. Here pre[4] - pre[1] = 10 - 4 = 6.

queries = [[0, 0], [1, 3], [0, 3]]
j = 0
while j < len(queries):
    left = queries[j][0]
    right = queries[j][1]
    print(pre[right + 1] - pre[left])
    j = j + 1

Output

4
6
10

Notice

One step up — with pre built once, EACH query is a single subtraction, so many queries cost one step each.

Put it together: the real program reads the whole input from stdin — n, the values, then the queries — builds pre once, and prints one total per query. The notebook has no contest input waiting, so run it from a terminal.

Try it yourself

import sys

data = sys.stdin.read()
tokens = data.split()
n = int(tokens[0])
query_count = int(tokens[1])
values = []
value_index = 0
while value_index < n:
    values.append(int(tokens[2 + value_index]))
    value_index = value_index + 1

pre = []
pre.append(0)
value_index = 0
while value_index < n:
    pre.append(pre[value_index] + values[value_index])
    value_index = value_index + 1

answer_text = ""
token_position = 2 + n
query_index = 0
while query_index < query_count:
    left = int(tokens[token_position])
    right = int(tokens[token_position + 1])
    range_total = pre[right + 1] - pre[left]
    if query_index > 0:
        answer_text = answer_text + "\n"
    answer_text = answer_text + str(range_total)
    token_position = token_position + 2
    query_index = query_index + 1
print(answer_text)

Run the full solver from this unit folder (the judge pipes a case to stdin):

python assets/l1.py < assets/l1/1.in

Notice

The full program is the same prefix idea — it just reads n, the values, and the queries from real stdin, builds pre once, and prints one total per line.

Test the edges: a single-element range (left == right), the full array (0 through n - 1), and a range ending at n - 1 — these expose mistakes around pre[0] and right + 1.

Complexity: building pre is O(n); each range query is O(1).

Lesson 2 — 2D Prefix Sums

Now a contest map asks for the point total inside many rectangles. Build a prefix grid with one extra zero row and one extra zero column; each cell stores its value plus the totals above and left, minus the upper-left overlap counted twice.

grid = [[2, 1, 4], [3, 5, 6]]
rows = 2
columns = 3
pre = []
r = 0
while r < rows + 1:
    prefix_row = []
    col = 0
    while col < columns + 1:
        prefix_row.append(0)
        col = col + 1
    pre.append(prefix_row)
    r = r + 1
r = 0
while r < rows:
    col = 0
    while col < columns:
        pre[r + 1][col + 1] = grid[r][col] + pre[r][col + 1] + pre[r + 1][col] - pre[r][col]
        col = col + 1
    r = r + 1
print(pre)

Output

[[0, 0, 0, 0], [0, 2, 3, 7], [0, 5, 11, 21]]

Notice

The prefix grid has an extra zero row and column; the build rule grid[r][c] + above + left - overlap gives rows [0,0,0,0], [0,2,3,7], [0,5,11,21].

r1 = 0
c1 = 1
r2 = 1
c2 = 2
print(pre[r2 + 1][c2 + 1] - pre[r1][c2 + 1] - pre[r2 + 1][c1] + pre[r1][c1])

Output

16

Notice

One step up — a rectangle total is four prefix corners: pre[r2+1][c2+1] - pre[r1][c2+1] - pre[r2+1][c1] + pre[r1][c1]. The rectangle (0,1)–(1,2) totals 21 - 0 - 5 + 0 = 16.

Put it together: the real program reads rows, columns, the grid, then the rectangle queries from stdin, builds the prefix grid once, and prints one total per query.

Try it yourself

import sys

data = sys.stdin.read()
tokens = data.split()
rows = int(tokens[0])
columns = int(tokens[1])
query_count = int(tokens[2])
grid = []
token_position = 3
row = 0
while row < rows:
    current_row = []
    column = 0
    while column < columns:
        current_row.append(int(tokens[token_position]))
        token_position = token_position + 1
        column = column + 1
    grid.append(current_row)
    row = row + 1

pre = []
row = 0
while row < rows + 1:
    prefix_row = []
    column = 0
    while column < columns + 1:
        prefix_row.append(0)
        column = column + 1
    pre.append(prefix_row)
    row = row + 1

row = 0
while row < rows:
    column = 0
    while column < columns:
        pre[row + 1][column + 1] = grid[row][column] + pre[row][column + 1] + pre[row + 1][column] - pre[row][column]
        column = column + 1
    row = row + 1

answer_text = ""
query_index = 0
while query_index < query_count:
    r1 = int(tokens[token_position])
    c1 = int(tokens[token_position + 1])
    r2 = int(tokens[token_position + 2])
    c2 = int(tokens[token_position + 3])
    rectangle_total = pre[r2 + 1][c2 + 1] - pre[r1][c2 + 1] - pre[r2 + 1][c1] + pre[r1][c1]
    if query_index > 0:
        answer_text = answer_text + "\n"
    answer_text = answer_text + str(rectangle_total)
    token_position = token_position + 4
    query_index = query_index + 1
print(answer_text)

Run the full solver from this unit folder (the judge pipes a case to stdin):

python assets/l2.py < assets/l2/1.in

Notice

The full program is the same inclusion-exclusion idea — it reads rows, columns, the grid, and the rectangle queries from real stdin, builds the prefix grid once, and prints one total per query.

Keep all four + 1 shifts and all four signs exactly; a missing corner changes rectangles that do not start at row 0, column 0. Test a single cell, the whole grid, and a rectangle reaching the last row and column.

Complexity: building the prefix grid is O(R * C); each rectangle query is O(1).

Recap

  • pre[0] = 0 and pre[i] = pre[i - 1] + a[i - 1], so pre has one more entry than a.
  • The total of a[left..right] is pre[right + 1] - pre[left].
  • Build pre once, then each query costs one step, however many queries there are.
  • A rectangle total is pre[r2+1][c2+1] - pre[r1][c2+1] - pre[r2+1][c1] + pre[r1][c1].
  • Test the edges: a single-element range, the whole array and a range ending at the last index.

You can now answer many range-sum questions fast with prefix sums.

Go to the exercises (9)

For parents and teachers: report a problem