RustArena

Step 7 of 7

Prefix sums

When you have to answer many "what's the total from day a to day b?" questions about data that doesn't change, precompute prefix sums: prefix[i] is the sum of the first i values, so prefix[0] = 0 and prefix has one more entry than the data. Then the sum of v[a..b] is just prefix[b] - prefix[a]: O(1) per question after O(n) setup.

fn main() {
    let sales = [120, 80, 45, 200, 15, 90, 60]; // one number per day
    let mut prefix = vec![0; sales.len() + 1];
    for i in 0..sales.len() {
        prefix[i + 1] = prefix[i] + sales[i];
    }
    println!("{prefix:?}");
    println!("days 2..5: {}", prefix[5] - prefix[2]);
    println!("whole week: {}", prefix[7] - prefix[0]);

    // The same running totals with an iterator
    let running: Vec<i32> = sales
        .iter()
        .scan(0, |total, &x| {
            *total += x;
            Some(*total)
        })
        .collect();
    println!("{running:?}");
}
[0, 120, 200, 245, 445, 460, 550, 610]
days 2..5: 260
whole week: 610
[120, 200, 245, 445, 460, 550, 610]

scan is an iterator adapter you haven't used yet: it carries a piece of state (total, starting at 0) from item to item and yields whatever the closure returns. It's handy for running totals.

The extra 0 at the front is what makes the formula work for ranges starting at day 0, with no special case.

Counting with prefix sums

Prefix sums also turn "a range with sum k" into "two prefixes that differ by k". While walking the data with a running total, the number of ranges ending here with sum k equals the number of earlier prefixes equal to total - k. Keep those counts in a HashMap and the whole count is O(n). Unlike the sliding window, this works with negative numbers.

The same idea works in two dimensions (sums over rectangles of an image or spreadsheet) with a 2D prefix table.

Common mistakes

Your task

Write a RangeSum type with fn new(values: &[i64]) -> Self (builds the prefix sums) and fn sum(&self, from: usize, to: usize) -> i64 (the sum of values[from..to], half-open, with from <= to <= len). The tests ask 200,000 questions about 200,000 values.

Challenge 2

Find the balance point of a list of weights: the first index i where the total of everything to the left of i equals the total of everything to the right of i (the item at i itself counts for neither side). Write fn balance_point(values: &[i64]) -> Option<usize>. For [1, 7, 3, 6, 5, 6] it's 3 (1 + 7 + 3 = 11 = 5 + 6).

Compute the total once. Then walk the list with a running left sum: the right side is total - left - values[i]. That's O(n); recomputing both sides at every index is O(n²), and the tests have 200,001 values.

Challenge 3

Finance wants to know how many runs of consecutive transactions net out to exactly k (for example, how many periods ended with no net change). Write fn count_subarrays_with_sum(values: &[i64], k: i64) -> u64. Values can be negative, so a sliding window won't do.

Walk the values with a running total and a HashMap<i64, u64> counting how often each prefix total has appeared, starting with the empty prefix (total 0, seen once). At each step, add the number of earlier prefixes equal to total - k, then record the current total. The tests include 100,000 values, where checking every range would be O(n²).

Previous: Sliding window