054. Count Subarrays Summing to K

Why this mechanism works

A range ending at boundary r sums to k exactly when an earlier prefix equals prefix[r]-k. Store frequencies of earlier prefix sums, because several boundaries may have the same sum. Seed the empty prefix with frequency one so ranges starting at index zero are counted.

State and boundaries

total is the current prefix sum. frequency contains only earlier boundary sums when queried. answer counts ranges ending at processed boundaries.

Protocol

1. Seed frequency[0]=1.
2. Extend total with the incoming value.
3. Add frequency[total-k] to answer.
4. Record the current total afterward.
5. Return answer.

Visualization

prefix sums for [1,-1,1]: 0,1,0,1
current 1 needs earlier 0 for k=1
there can be more than one earlier match

Python implementation

def count_sum_k(nums, k):
    frequency = {0: 1}
    total = answer = 0
    for value in nums:
        total += value
        answer += frequency.get(total - k, 0)
        frequency[total] = frequency.get(total, 0) + 1
    return answer

Worked execution

Recorded state changes after statements; L numbers refer to the solution file. Long traces are explicitly condensed.

L2 count_sum_k: frequency = {0: 1}
nums=[1, -1, 1]; k=1; frequency={0:1}

L3 count_sum_k: total = answer = 0
nums=[1, -1, 1]; k=1; frequency={0:1}; total=0; answer=0

L4 count_sum_k: for value in nums:
nums=[1, -1, 1]; k=1; frequency={0:1}; total=0; answer=0; value=1

L5 count_sum_k: total += value
nums=[1, -1, 1]; k=1; frequency={0:1}; total=1; answer=0; value=1

L6 count_sum_k: answer += frequency.get(total - k, 0)
nums=[1, -1, 1]; k=1; frequency={0:1}; total=1; answer=1; value=1

L7 count_sum_k: frequency[total] = frequency.get(total, 0) + 1
nums=[1, -1, 1]; k=1; frequency={0:1, 1:1}; total=1; answer=1; value=1

L4 count_sum_k: for value in nums:
nums=[1, -1, 1]; k=1; frequency={0:1, 1:1}; total=1; answer=1; value=-1

L5 count_sum_k: total += value
nums=[1, -1, 1]; k=1; frequency={0:1, 1:1}; total=0; answer=1; value=-1

... 1 intermediate state changes omitted ...

L4 count_sum_k: for value in nums:
nums=[1, -1, 1]; k=1; frequency={0:2, 1:1}; total=0; answer=1; value=1

L5 count_sum_k: total += value
nums=[1, -1, 1]; k=1; frequency={0:2, 1:1}; total=1; answer=1; value=1

L6 count_sum_k: answer += frequency.get(total - k, 0)
nums=[1, -1, 1]; k=1; frequency={0:2, 1:1}; total=1; answer=3; value=1

L7 count_sum_k: frequency[total] = frequency.get(total, 0) + 1
nums=[1, -1, 1]; k=1; frequency={0:2, 1:2}; total=1; answer=3; value=1

L2 count_sum_k: frequency = {0: 1}
nums=[0, 0]; k=0; frequency={0:1}

L3 count_sum_k: total = answer = 0
nums=[0, 0]; k=0; frequency={0:1}; total=0; answer=0

L4 count_sum_k: for value in nums:
nums=[0, 0]; k=0; frequency={0:1}; total=0; answer=0; value=0

L6 count_sum_k: answer += frequency.get(total - k, 0)
nums=[0, 0]; k=0; frequency={0:1}; total=0; answer=1; value=0

L7 count_sum_k: frequency[total] = frequency.get(total, 0) + 1
nums=[0, 0]; k=0; frequency={0:2}; total=0; answer=1; value=0

L6 count_sum_k: answer += frequency.get(total - k, 0)
nums=[0, 0]; k=0; frequency={0:2}; total=0; answer=3; value=0

L7 count_sum_k: frequency[total] = frequency.get(total, 0) + 1
nums=[0, 0]; k=0; frequency={0:3}; total=0; answer=3; value=0

Why the solution is correct

Each earlier matching prefix identifies exactly one nonempty interval ending here. Querying before insertion prevents using the current boundary as both endpoints. Each interval has one right boundary, so it is counted exactly once.

Time and space

O(n) expected time and O(n) auxiliary space.

Common mistakes

A set loses multiple equal-prefix boundaries. Inserting before querying overcounts empty ranges when k=0.

Reasoning card

Transfer exercise

Count sums divisible by a positive modulus by grouping prefix sums with equal remainders.