Prefix Sum Technique in Python: Range Sum, Subarray Sum Equals K and 2D Prefix Sums

Suppose you must answer the same question hundreds of times: what is the sum between index 3 and index 8? Adding the numbers each time repeats work. Prefix sums do it once up front, and every later query becomes a subtraction. The same trick solves counting problems such as Subarray Sum Equals K, even with negative numbers. Below are the classic problems with tested code, plus 2D prefix sums and difference arrays.

What is a prefix sum?

A prefix sum array stores, at each position, the total of all elements before it. If prefix[i] is the sum of the first i elements, then the sum of the elements from index left to index right (inclusive) is:

prefix[right + 1] - prefix[left]

The reasoning is simple. prefix[right + 1] contains everything up to right. Subtracting prefix[left] removes everything before left, leaving only the range you want. Building the table costs O(n) once. Each query then costs O(1), instead of O(n) for adding the numbers again.

Problem 1: Range Sum Query

Given an array, answer many queries of the form: what is the sum between indexes left and right?

class NumArray:
    def __init__(self, nums):
        self.prefix = [0]                        # prefix[0] = 0 on purpose
        for n in nums:
            self.prefix.append(self.prefix[-1] + n)

    def sum_range(self, left, right):
        return self.prefix[right + 1] - self.prefix[left]

arr = NumArray([-2, 0, 3, -5, 2, -1])
print(arr.sum_range(0, 2))   # 1
print(arr.sum_range(2, 5))   # -1
print(arr.sum_range(0, 5))   # -3

For [-2, 0, 3, -5, 2, -1] the prefix array is [0, -2, -2, 1, -4, -2, -3]. To get the sum from index 2 to 5, compute prefix[6] - prefix[2], which is -3 – (-2) = -1. Notice the leading 0. It represents the empty prefix and lets the formula work even when left is 0, so you never need a special case. With q queries, the total time is O(n + q) instead of O(n x q).

Problem 2: Subarray Sum Equals K

Count the contiguous subarrays whose sum equals k. This is where prefix sums earn their place, because the array may contain negative numbers.

Write any subarray sum as a difference of two prefix sums: the subarray that ends at the current position has sum current_prefix - earlier_prefix. We want that to equal k, which means earlier_prefix = current_prefix - k. So at each position we only need to know how many earlier prefix sums had that exact value. A hash map from prefix sum to its count answers that instantly.

def subarray_sum(nums, k):
    seen = {0: 1}               # the empty prefix has sum 0, seen once
    total = count = 0
    for n in nums:
        total += n
        count += seen.get(total - k, 0)
        seen[total] = seen.get(total, 0) + 1
    return count

print(subarray_sum([1, 1, 1], 2))    # 2
print(subarray_sum([1, 2, 3], 3))    # 2
print(subarray_sum([1, -1, 0], 0))   # 3

Trace on [1, 2, 3] with k = 3

1. Start with seen = {0: 1}, total = 0, count = 0.
2. n = 1. total = 1. We look for 1 – 3 = -2, which is not present. Store {0: 1, 1: 1}.
3. n = 2. total = 3. We look for 3 – 3 = 0, which is present once, so count becomes 1. That is the subarray [1, 2]. Store 3.
4. n = 3. total = 6. We look for 6 – 3 = 3, which is present once, so count becomes 2. That is the subarray [3].

The answer is 2. Three details matter. First, the map starts with {0: 1} so that a subarray starting at index 0 is counted. Second, we update the map after checking, so a subarray is never empty. Third, we need a count map, not a set, because the same prefix sum can appear several times and each occurrence is a different starting point. Time is O(n) and space is O(n).

Why not a sliding window?

A sliding window works when adding an element always increases the sum, which needs positive numbers. With negative numbers, a longer window can have a smaller sum, so you cannot decide when to shrink. Prefix sums with a hash map do not depend on the sign of the numbers. If your input is all positive and you only need the longest or shortest window, the sliding window pattern is simpler and uses less memory.

Problem 3: longest subarray with sum k

A small change gives a different answer. To find the length of the longest subarray with sum k, store the first index at which each prefix sum appears. The earliest start gives the longest subarray.

def longest_subarray_sum_k(nums, k):
    first = {0: -1}             # prefix sum -> earliest index
    total = best = 0
    for i, n in enumerate(nums):
        total += n
        if total - k in first:
            best = max(best, i - first[total - k])
        first.setdefault(total, i)
    return best

print(longest_subarray_sum_k([1, -1, 5, -2, 3], 3))   # 4

For [1, -1, 5, -2, 3] and k = 3, the subarray [1, -1, 5, -2] sums to 3 and has length 4. We use setdefault so that a later repeat of the same prefix sum does not overwrite the earlier, better index.

Problem 4: Product of Array Except Self

Prefix ideas work with products too. Given an array, return a new array where each element is the product of all other elements, without using division and in O(n) time. For each position, the answer is the product of everything to its left times the product of everything to its right.

def product_except_self(nums):
    n = len(nums)
    res = [1] * n
    left = 1
    for i in range(n):
        res[i] = left            # product of everything to the left
        left *= nums[i]
    right = 1
    for i in range(n - 1, -1, -1):
        res[i] *= right          # multiply by product of everything to the right
        right *= nums[i]
    return res

print(product_except_self([1, 2, 3, 4]))      # [24, 12, 8, 6]
print(product_except_self([-1, 1, 0, -3, 3]))  # [0, 0, 9, 0, 0]

For [1, 2, 3, 4], the first pass stores the left products [1, 1, 2, 6]. The second pass walks from the right, multiplying by 4, then 12, then 24, then 24, which produces [24, 12, 8, 6]. Reusing the result array for the left products keeps the extra space at O(1), not counting the output.

Problem 5: 2D prefix sums

The same idea extends to grids. Let p[r][c] be the sum of the rectangle from the top-left corner to just before row r and column c. Each cell is built from three neighbours using inclusion and exclusion: add the cell above and the cell to the left, then subtract the overlap that was counted twice.

class NumMatrix:
    def __init__(self, matrix):
        rows, cols = len(matrix), len(matrix[0])
        self.p = [[0] * (cols + 1) for _ in range(rows + 1)]
        for r in range(rows):
            for c in range(cols):
                self.p[r + 1][c + 1] = (matrix[r][c] + self.p[r][c + 1]
                                        + self.p[r + 1][c] - self.p[r][c])

    def sum_region(self, r1, c1, r2, c2):
        p = self.p
        return p[r2 + 1][c2 + 1] - p[r1][c2 + 1] - p[r2 + 1][c1] + p[r1][c1]

To get the sum of any rectangle, take the big rectangle that reaches the bottom-right corner, subtract the strip above and the strip to the left, and add back the corner that was subtracted twice. Building takes O(rows x cols), and each query takes O(1). Be careful with the extra row and column of zeros, which play the same role as the leading 0 in the one-dimensional version.

Bonus: the difference array for range updates

Prefix sums answer many range questions. Their mirror image, the difference array, handles many range updates. To add a value v to every element from l to r, add v at position l and subtract v at position r + 1. When you finally take the running total, each update has applied to exactly the right range.

def apply_updates(length, updates):
    diff = [0] * (length + 1)
    for l, r, v in updates:
        diff[l] += v
        diff[r + 1] -= v
    res, cur = [], 0
    for i in range(length):
        cur += diff[i]
        res.append(cur)
    return res

print(apply_updates(5, [(1, 3, 2), (2, 4, 3), (0, 2, -2)]))
# [-2, 0, 3, 5, 3]

Each update costs O(1), so q updates on an array of length n cost O(n + q) in total. This idea appears in problems such as Corporate Flight Bookings and Car Pooling.

Complexity summary

If you want to practise explaining these costs clearly, our Big O guide covers the vocabulary interviewers expect.

How to recognise a prefix sum problem

• You must answer many range-sum queries on a fixed array.
• The question counts or finds subarrays with a given sum, especially when negatives are allowed.
• Each answer depends on everything to the left and everything to the right of a position.
• You need many range updates followed by one final read of the array.
• A brute force solution recomputes the same sum over and over.

Common mistakes

• Off-by-one in the formula. With a leading zero, the sum from left to right is prefix[right + 1] - prefix[left]. Forgetting the +1 is the most common bug.
• Forgetting the seed entry. Without {0: 1} in the hash map, subarrays that start at index 0 are missed.
• Updating the map before the lookup. This lets an empty subarray be counted when k is 0.
• Using a set instead of a count map. Repeated prefix sums then undercount the answer.
• Using sliding window with negatives. The window logic breaks, so use prefix sums.

Practice problems

• Range Sum Query – Immutable
• Subarray Sum Equals K
• Contiguous Array and Maximum Size Subarray Sum Equals k
• Product of Array Except Self
• Range Sum Query 2D – Immutable
• Car Pooling (difference array)

Python’s itertools.accumulate builds a prefix array in one line, and you can read more about it in our Python cheat sheet for coding interviews. For the complement idea behind the hash map step, revisit Two Sum.

FAQ

Final checklist

A leading zero, prefix[right + 1] - prefix[left], a hash map seeded with zero, and a count instead of a set. If those four details are right, most prefix sum problems fall in a few lines.

Keep learning: sliding window, Two Sum and monotonic stack build on the ideas in this guide.

Also see: interval problems for related patterns.