The Running Total

An easy Python interview practice problem on DataDriven. Write the Python and run it against test cases, with instant feedback.

Domain
Python
Difficulty
easy
Seniority
L3

Problem

Return our account's balance after each day, given nums, the ledger's daily balance changes in day order, where a withdrawal arrives as a negative change. An empty ledger yields an empty list of balances.

Each position holds the sum of everything before it.

Example

Input:

nums = [1,2,3,4,5]

Output:

[1,3,6,10,15]

Worked solution and explanation

What this really is

This is a running balance, and the whole question is whether you keep every intermediate value. The tempting read of 'sum the ledger' is one number: sum(nums) hands back 15 for the example, in one line. The grader wants the balance after every day. The answer has exactly as many entries as nums, the first is day one's change on its own, and an empty ledger comes back as an empty list rather than 0. Return the scalar and every case fails on shape before the values even matter.

Step 1: Start the balance at zero

Before day one nothing has happened, so total = 0. That is also what makes the empty ledger work without extra code: the loop never runs and the empty balances list goes back untouched.

Step 2: Fold in each change, then record it

Add the day's change to total and append the new total right away. The order matters: append before adding and every balance is one day late, so the example would read 0, 1, 3, 6, 10.

Step 3: Let negatives do nothing special

A withdrawal is just a negative change, and total += change already pulls the balance down. On [-2, 5, -3] the balances are -2, 3, 0. A zero balance is a real reading and belongs in the output.

The solution

Carry the balance forward, one entry per day
def prefix_sum(nums):
    balances = []
    total = 0
    for change in nums:
        total += change
        balances.append(total)
    return balances
Recompute each day
sum(nums[:i + 1]) for every day gives the right numbers, but each day re-adds the whole history. A year of entries costs about 67,000 additions instead of 365.
Carry the balance
A single total carried forward adds each change exactly once, so the pass is O(n). The only extra memory is the output you were asked for.

Common follow-up questions