Go Guide

Core patterns · Lesson 14 of 62 · about 45 min

Arrays & prefix sums

Precompute running totals once, then answer any range-sum question with one subtraction, and count subarrays with a hash map.

Go by Example: Slices ↗
On this page · 9 sections
  1. 01Building the prefix array
  2. 02Answering a range query
  3. 03Counting subarrays with a given sum
  4. 04The full program
  5. 05Complexity
  6. 06Recognizing the pattern
  7. 07Summary
  8. 08Check your understanding
  9. 09Practice in Go

Building the prefix array#

Define pre[i] as the sum of the first i elements. Make it one longer than the input so pre[0] = 0:

Go
pre := make([]int, len(nums)+1)
for i, n := range nums {
	pre[i+1] = pre[i] + n
}

Trace it for nums = [3, 1, 4, 1, 5]:

inums[i]pre[i+1] = pre[i] + nums[i]
pre[0] = 0
03pre[1] = 0 + 3 = 3
11pre[2] = 3 + 1 = 4
24pre[3] = 4 + 4 = 8
31pre[4] = 8 + 1 = 9
45pre[5] = 9 + 5 = 14

Answering a range query#

The sum of nums[l..r] (both inclusive) is everything up to r minus everything before l:

sum(l, r) = pre[r+1] - pre[l]

For l = 1, r = 3: pre[4] - pre[1] = 9 - 3 = 6, and indeed 1 + 4 + 1 = 6.

Picture it as two bars: pre[r+1] covers indexes 0..r, pre[l] covers 0..l-1; subtract and you are left with exactly l..r.

Run the trace

This program prints the same table as above while it builds pre, then answers several range queries from that one array:

trace.go
Go
package main
 
import "fmt"
 
func main() {
	nums := []int{3, 1, 4, 1, 5}
	pre := make([]int, len(nums)+1)
	fmt.Println("pre[0] = 0")
	for i, n := range nums {
		pre[i+1] = pre[i] + n
		fmt.Printf("i=%d nums[i]=%d  pre[%d] = %d + %d = %d\n", i, n, i+1, pre[i], n, pre[i+1])
	}
	fmt.Println("pre =", pre)
 
	// Every query is one subtraction, no matter how long the range is.
	for _, q := range [][2]int{{1, 3}, {0, 4}, {2, 2}, {3, 4}} {
		l, r := q[0], q[1]
		fmt.Printf("sum(%d,%d) = pre[%d] - pre[%d] = %2d - %d = %d\n",
			l, r, r+1, l, pre[r+1], pre[l], pre[r+1]-pre[l])
	}
}
Output
pre[0] = 0
i=0 nums[i]=3  pre[1] = 0 + 3 = 3
i=1 nums[i]=1  pre[2] = 3 + 1 = 4
i=2 nums[i]=4  pre[3] = 4 + 4 = 8
i=3 nums[i]=1  pre[4] = 8 + 1 = 9
i=4 nums[i]=5  pre[5] = 9 + 5 = 14
pre = [0 3 4 8 9 14]
sum(1,3) = pre[4] - pre[1] =  9 - 3 = 6
sum(0,4) = pre[5] - pre[0] = 14 - 0 = 14
sum(2,2) = pre[3] - pre[2] =  8 - 4 = 4
sum(3,4) = pre[5] - pre[3] = 14 - 8 = 6

Change nums and predict the pre line before you run it again.

Step through it

Try your own numbers and query range.

Counting subarrays with a given sum#

"How many contiguous subarrays add up to k?" Brute force tries every start and end: O(n²). With prefix sums:

  • A subarray from i+1 to j sums to pre[j] - pre[i].
  • We want pre[j] - pre[i] = k, which means pre[i] = pre[j] - k.
  • So as we walk with a running sum, we ask: "how many earlier prefix sums equal sum - k?" A map from sum to count answers in O(1).

Trace nums = [1, 2, 1, 2, 1], k = 3:

nrunning sumlook up sum - kcount addedmap after
0{0:1}
11-2 → 00{0:1, 1:1}
230 → 11 ([1,2]){0:1, 1:1, 3:1}
141 → 11 ([2,1]){…, 4:1}
263 → 11 ([1,2]){…, 6:1}
174 → 11 ([2,1]){…, 7:1}

Total: 4. The {0:1} seed is what lets the first [1,2] count, because its "earlier prefix" is the empty prefix.

Run the trace

The program prints one line per row of the table. At the end, a brute-force double loop lists the matching subarrays so you can see that the map counted exactly the right ones:

count_trace.go
Go
package main
 
import "fmt"
 
func main() {
	nums, k := []int{1, 2, 1, 2, 1}, 3
	seen := map[int]int{0: 1} // the empty prefix
	sum, count := 0, 0
	for j, n := range nums {
		sum += n
		found := seen[sum-k] // earlier prefixes that leave exactly k
		count += found
		fmt.Printf("j=%d n=%d sum=%d  need %2d: seen %d time(s)  count=%d\n", j, n, sum, sum-k, found, count)
		seen[sum]++
	}
 
	fmt.Println("brute force check:")
	for i := range nums {
		s := 0
		for j := i; j < len(nums); j++ {
			s += nums[j]
			if s == k {
				fmt.Printf("  nums[%d:%d] = %v\n", i, j+1, nums[i:j+1])
			}
		}
	}
}
Output
j=0 n=1 sum=1  need -2: seen 0 time(s)  count=0
j=1 n=2 sum=3  need  0: seen 1 time(s)  count=1
j=2 n=1 sum=4  need  1: seen 1 time(s)  count=2
j=3 n=2 sum=6  need  3: seen 1 time(s)  count=3
j=4 n=1 sum=7  need  4: seen 1 time(s)  count=4
brute force check:
  nums[0:2] = [1 2]
  nums[1:3] = [2 1]
  nums[2:4] = [1 2]
  nums[3:5] = [2 1]

The full program#

prefix.go
Go
package main
 
import "fmt"
 
type NumArray struct{ pre []int }
 
func NewNumArray(nums []int) *NumArray {
	pre := make([]int, len(nums)+1)
	for i, n := range nums {
		pre[i+1] = pre[i] + n
	}
	return &NumArray{pre}
}
 
// SumRange returns nums[l] + ... + nums[r].
func (a *NumArray) SumRange(l, r int) int { return a.pre[r+1] - a.pre[l] }
 
// subarraySum counts contiguous subarrays summing to k in O(n).
func subarraySum(nums []int, k int) int {
	count, sum := 0, 0
	seen := map[int]int{0: 1}
	for _, n := range nums {
		sum += n
		count += seen[sum-k]
		seen[sum]++
	}
	return count
}
 
// productExceptSelf uses a prefix pass and a suffix pass, no division.
func productExceptSelf(nums []int) []int {
	out := make([]int, len(nums))
	left := 1
	for i, n := range nums {
		out[i] = left
		left *= n
	}
	right := 1
	for i := len(nums) - 1; i >= 0; i-- {
		out[i] *= right
		right *= nums[i]
	}
	return out
}
 
func main() {
	a := NewNumArray([]int{3, 1, 4, 1, 5})
	fmt.Println(a.pre, a.SumRange(1, 3), a.SumRange(0, 4), a.SumRange(2, 2))
	fmt.Println(subarraySum([]int{1, 2, 1, 2, 1}, 3), subarraySum([]int{1, -1, 0}, 0))
	fmt.Println(productExceptSelf([]int{1, 2, 3, 4}))
}
Output
[0 3 4 8 9 14] 6 14 4
4 3
[24 12 8 6]

Complexity#

OperationTimeSpace
build prefix arrayO(n)O(n)
one range queryO(1)
count subarrays with sum kO(n)O(n) map
product except selfO(n)O(1) besides the output

Recognizing the pattern#

Reach for prefix sums when you see: "sum of a range", "many queries on a fixed array", "subarray with sum equal to / divisible by k", or "balance point" (left side equals right side). The same idea works with products, XOR, counts of a character, and in two dimensions for grids.

Prefix sums in two dimensions

For a grid, let pre[r][c] be the sum of the rectangle from the top-left corner up to (but not including) row r and column c. Each cell is: the value, plus the rectangle above it, plus the rectangle to its left, minus their overlap, which was added twice. Any sub-rectangle then takes four lookups:

prefix2d.go
Go
package main
 
import "fmt"
 
func build(grid [][]int) [][]int {
	pre := make([][]int, len(grid)+1)
	for i := range pre {
		pre[i] = make([]int, len(grid[0])+1)
	}
	for r, row := range grid {
		for c, v := range row {
			pre[r+1][c+1] = v + pre[r][c+1] + pre[r+1][c] - pre[r][c]
		}
	}
	return pre
}
 
// rect sums grid[r1..r2][c1..c2], corners included.
func rect(pre [][]int, r1, c1, r2, c2 int) int {
	return pre[r2+1][c2+1] - pre[r1][c2+1] - pre[r2+1][c1] + pre[r1][c1]
}
 
func main() {
	grid := [][]int{
		{1, 2, 3, 4},
		{5, 6, 7, 8},
		{9, 10, 11, 12},
	}
	pre := build(grid)
	for _, row := range pre {
		fmt.Println(row)
	}
	fmt.Println("middle 2x2:", rect(pre, 1, 1, 2, 2)) // 6 + 7 + 10 + 11
	fmt.Println("whole grid:", rect(pre, 0, 0, 2, 3))
	fmt.Println("one cell:  ", rect(pre, 2, 0, 2, 0))
}
Output
[0 0 0 0 0]
[0 1 3 6 10]
[0 6 14 24 36]
[0 15 33 54 78]
middle 2x2: 34
whole grid: 78
one cell:   9

Summary#

  • pre[i] is the sum of the first i elements, with pre[0] = 0.
  • Range sum l..r is pre[r+1] - pre[l], in O(1) after O(n) setup.
  • Prefix sums plus a hash map count subarrays with a target sum in O(n), negatives included.
  • The idea extends to products, XOR and 2D grids.

Check your understanding

Answer before scrolling back up. If you miss one, reread the matching section, then try again.

  1. Q1.With pre of length n+1 (pre[0] = 0), the sum of nums[2..5] inclusive is:

  2. Q2.Why make the prefix array one longer than the input?

  3. Q3.In Subarray Sum Equals K, why seed the map with {0: 1}?

  4. Q4.Building a prefix array costs O(n). Answering q range queries afterwards costs:

Practice in Go

Write a table test first, then the solution. On LeetCode, choose Go as the language. Open a hint only after 15 minutes stuck.

Finished this lesson?

Mark it done to track your progress. It is saved in this browser only.