Solving Range Queries with Prefix Sums

Welcome to the next lesson in our journey through essential algorithmic patterns. In our previous session, we mastered techniques for handling interval-based problems, which are common when dealing with schedules or resource allocation. We saw how sorting is often the key to unlocking an efficient solution.

Today, we'll explore another powerful technique for problems involving ranges of data: the prefix sum. Imagine you're analyzing performance metrics over time, processing financial transactions, or working with any cumulative data. You'll often face "range query" problems, where you need to calculate a sum or an aggregate over a specific sub-segment of data, and you need to do it fast. The prefix sum technique provides an elegant way to answer these queries in constant time after a one-time preprocessing step.

By the end of this lesson, you will be able to apply this technique to efficiently solve range query problems in both one and two dimensions.

The Inefficiency of Brute-Force Range Queries

Let's start with a classic problem: you are given an array of numbers, and you need to handle numerous queries, each asking for the sum of elements within a specific range [L, R].

The most straightforward approach is to loop from the left index L to the right index R for each query and add up the elements.

// Naive approach
func rangeSum(nums []int, L int, R int) int {
    sum := 0
    for i := L; i <= R; i++ {
        sum += nums[i]
    }
    return sum
}

If the array has N elements and you receive Q queries, the time complexity of this approach is $O(N\cdot Q)$ in the worst case (when each query spans the whole array). If N and Q are large, say $10^5$, this becomes $10^{10}$ operations, which is far too slow for any real-world application or coding interview. This is a clear signal that we need a more efficient method.

The Prefix Sum Solution in One Dimension

The prefix sum technique is all about trading a little bit of upfront computation and space for incredibly fast queries. The core idea is to create a new array, often called prefixSum, where each element prefixSum[i] stores the cumulative sum of the original array's elements from index 0 up to i.

Mathematically, for an original array A:

$$\text{prefixSum}[i]=A[0]+A[1]+\cdots +A[i]$$

The following video provides an excellent, intuitive explanation of this concept, using an analogy of YouTube video performance to illustrate why pre-computation is so powerful.

As the video explained, once we have the prefix sum array, calculating the sum of any range [L, R] becomes a simple subtraction that takes $O(1)$ time.

The sum of elements from index L to R is the total sum up to R minus the total sum up to L-1.

• Sum(A[L...R]) = prefixSum[R] - prefixSum[L-1]

The image below visualizes this exact calculation.

A small but important detail is handling the edge case where L = 0. The formula prefixSum[R] - prefixSum[L-1] would try to access prefixSum[-1]. We can handle this with a simple if condition or by using the padding trick shown in the video, where an extra 0 is added to the start of the prefix sum array.

Here is a clean implementation in Go that uses a simple conditional check.

package main

import "fmt"

// NumArray holds the pre-computed prefix sums.
type NumArray struct {
	prefixSums []int
}

// Constructor creates the NumArray object and calculates prefix sums.
func Constructor(nums []int) NumArray {
	if len(nums) == 0 {
		return NumArray{prefixSums: []int{}}
	}

	// O(N) preprocessing step
	prefixSums := make([]int, len(nums))
	prefixSums[0] = nums[0]
	for i := 1; i < len(nums); i++ {
		prefixSums[i] = prefixSums[i-1] + nums[i]
	}

	return NumArray{prefixSums: prefixSums}
}

// SumRange answers a query in O(1) time.
func (na *NumArray) SumRange(left int, right int) int {
	if left == 0 {
		return na.prefixSums[right]
	}
	return na.prefixSums[right] - na.prefixSums[left-1]
}

func main() {
	nums := []int{-2, 0, 3, -5, 2, -1}
	obj := Constructor(nums)
	fmt.Println(obj.SumRange(0, 2)) // Output: 1
	fmt.Println(obj.SumRange(2, 5)) // Output: -1
	fmt.Println(obj.SumRange(0, 5)) // Output: -3
}

The overall time complexity is now $O(N+Q)$, a massive improvement over $O(N\cdot Q)$. The space complexity is $O(N)$ to store the prefix sum array.

Extending to 2D: Submatrix Sum Queries

The prefix sum pattern elegantly extends to two-dimensional arrays (matrices). The problem is analogous: given a matrix, efficiently calculate the sum of elements within a given rectangular subregion defined by its top-left (r1, c1) and bottom-right (r2, c2) corners.

To build the 2D prefix sum matrix P, we use the principle of inclusion-exclusion. The sum of the rectangle ending at (i, j) is:

$$P[i][j]=A[i][j]+P[i-1][j]+P[i][j-1]-P[i-1][j-1]$$

We subtract P[i-1][j-1] because its area was included twice (once in P[i-1][j] and once in P[i][j-1]).

Querying also uses the same principle. The sum of the submatrix from (r1, c1) to (r2, c2) is:

$$\text{Sum}=P[r2][c2]-P[r1-1][c2]-P[r2][c1-1]+P[r1-1][c1-1]$$

Again, we add back P[r1-1][c1-1] because it was subtracted twice.

The following video provides a brilliant walkthrough of this 2D logic.

Here is a Go implementation for 2D range sum queries. Note the use of a padded matrix to simplify the boundary logic, similar to the trick for the 1D case.

package main

import "fmt"

type NumMatrix struct {
	prefixSum2D [][]int
}

func Constructor2D(matrix [][]int) NumMatrix {
	if len(matrix) == 0 || len(matrix[0]) == 0 {
		return NumMatrix{}
	}
	rows, cols := len(matrix), len(matrix[0])
	
    // Create a padded prefix sum matrix to avoid bounds checks.
	// Size is (rows+1) x (cols+1)
	prefixSum2D := make([][]int, rows+1)
	for i := range prefixSum2D {
		prefixSum2D[i] = make([]int, cols+1)
	}

	// O(M*N) preprocessing
	for r := 1; r <= rows; r++ {
		for c := 1; c <= cols; c++ {
			prefixSum2D[r][c] = matrix[r-1][c-1] +
				prefixSum2D[r-1][c] +
				prefixSum2D[r][c-1] -
				prefixSum2D[r-1][c-1]
		}
	}
	return NumMatrix{prefixSum2D: prefixSum2D}
}

// O(1) query
func (nm *NumMatrix) SumRegion(r1, c1, r2, c2 int) int {
	// Adjust coordinates for the padded matrix
	r1, c1, r2, c2 = r1+1, c1+1, r2+1, c2+1
	
	return nm.prefixSum2D[r2][c2] -
		nm.prefixSum2D[r1-1][c2] -
		nm.prefixSum2D[r2][c1-1] +
		nm.prefixSum2D[r1-1][c1-1]
}

func main() {
	matrix := [][]int{
		{3, 0, 1, 4, 2},
		{5, 6, 3, 2, 1},
		{1, 2, 0, 1, 5},
		{4, 1, 0, 1, 7},
		{1, 0, 3, 0, 5},
	}
	obj := Constructor2D(matrix)
	fmt.Println(obj.SumRegion(2, 1, 4, 3)) // Output: 8
	fmt.Println(obj.SumRegion(1, 1, 2, 2)) // Output: 11
}

A Key Application: Subarray Sum Equals K

The prefix sum pattern is not just for simple range sums. It's a foundational tool for solving more complex problems. One of the most famous is "Subarray Sum Equals K": given an array and an integer K, find the total number of subarrays whose elements sum to K.

This problem beautifully combines prefix sums with the hash map pattern we studied earlier. Let's break it down:

1. We know that the sum of a subarray from index i to j is prefixSum[j] - prefixSum[i-1].
2. We want this sum to be equal to K: prefixSum[j] - prefixSum[i-1] = K.
3. Rearranging the equation, we get: prefixSum[j] - K = prefixSum[i-1].

This is the crucial insight. As we iterate through the array and calculate the current prefix sum (currentSum = prefixSum[j]), we need to find out how many times we have previously seen a prefix sum equal to currentSum - K.

A hash map is the perfect tool for this. We can store the frequencies of all prefix sums we've encountered so far.

The following segment of the video you watched earlier explains this exact logic with exceptional clarity.

Here's a Go implementation that puts this idea into practice.

package main

import "fmt"

func subarraySum(nums []int, k int) int {
	count := 0
	currentSum := 0
	
	// A map to store the frequency of prefix sums encountered so far.
	// Key: prefix sum, Value: frequency.
	prefixSumFreq := make(map[int]int)
	
	// Base case: a prefix sum of 0 has been seen once (the empty prefix).
	// This handles cases where a subarray starting from index 0 sums to k.
	prefixSumFreq[0] = 1
	
	for _, num := range nums {
		currentSum += num
		
		// We are looking for a previous prefix sum `p` such that:
		// currentSum - p = k  =>  p = currentSum - k
		neededSum := currentSum - k
		
		// If we've seen this neededSum before, it means we found valid subarrays.
		if freq, found := prefixSumFreq[neededSum]; found {
			count += freq
		}
		
		// Update the frequency of the current prefix sum in the map.
		prefixSumFreq[currentSum]++
	}
	
	return count
}

func main() {
	nums := []int{1, 1, 1}
	k := 2
	fmt.Println(subarraySum(nums, k)) // Output: 2 ([1,1] starting at index 0 and [1,1] starting at index 1)
	
	nums2 := []int{1, 2, 3}
	k2 := 3
	fmt.Println(subarraySum(nums2, k2)) // Output: 2 ([1,2] and [3])
}

Conclusion

In this lesson, we've explored the prefix sum technique, a fundamental pattern for efficiently handling range-based problems. By investing in an $O(N)$ pre-computation step, we unlock the ability to answer countless queries in $O(1)$ time.

Key Takeaways:

• Pre-computation for Efficiency: The core idea is to trade space and a linear-time setup for constant-time queries, which is a huge win when the number of queries is large.
• 1D and 2D Formulas: You learned the simple subtraction formula for 1D range sums and the inclusion-exclusion principle for 2D submatrix sums.
• Versatility of the Pattern: The prefix sum is not just for summing. It's a building block for more complex problems, as demonstrated with the "Subarray Sum Equals K" problem, where it pairs perfectly with hash maps.

You've now added another essential tool to your algorithmic arsenal. In our next lesson, we'll transition to a new data structure and module, starting with Linked Lists. We'll begin by tackling a classic interview problem: the in-place reversal of a singly linked list.