Prefix Sum

The prefix sum is an incredibly powerful and straightforward technique. Its primary goal is to allow for constant-time range sum queries on an array.

What is Prefix Sum?

The prefix sum of an array at index 'i' is the sum of all numbers from index '0' to 'i'. By computing and storing the prefix sum of an array, you can easily answer queries about the sum of any subarray in constant time.

# Example for Prefix Sum in Python

def prefix_sum_array(arr):
    prefix_sum = [0]
    for num in arr:
        prefix_sum.append(prefix_sum[-1] + num)
    return prefix_sum

arr = [1, -20, -3, 30, 5, 4]
print(prefix_sum_array(arr))

Prefix Sum vs. Sliding Window

While both techniques solve subarray-related problems, they have different use cases:

  • Prefix Sum: Ideal for quick range sum queries on static arrays.
  • Sliding Window: Better for problems involving contiguous subarrays with specific conditions or sizes, e.g., maximum sum of subarrays of size 'k'.

Common Subarray Problems

  1. Subarray with Given Sum: Find a contiguous subarray that sums to a target value.
  2. Zero Sum Subarrays: Locate subarrays that sum to zero.
  3. Longest Subarray with Sum k: Find the longest subarray with a specific sum.
  4. Smallest Subarray with Sum > X: Identify the smallest subarray whose sum exceeds a given value.
  5. Maximum Sum of Subarrays of Size k: Find the maximum sum among all subarrays of a fixed size.
  6. Suffix Sum: Similar to prefix sum, but calculated from the end of the array.

Suffix Sum

Just as we have a prefix sum that calculates the sum from the beginning to a specific index, we can also calculate a suffix sum, which calculates the sum from a particular index to the end of the array.

Now let's practice one of such problems.

Practice Problem: Subarray Sum

Given an array of integers and an integer target, find a subarray that sums to target and return the start and end indices of the subarray.

Input: arr: 1 -20 -3 30 5 4 target: 7

Output: 1 4

Explanation: -20 - 3 + 30 = 7. The indices for subarray [-20,-3,30] is 1 and 4 (right exclusive).

Try it yourself

Explanation

Intuition

The brute force way is to find the sum of each subarray and compare it with the target. Let N be the number of elements in the array, there are N subarrays with size 1, N-1 subarrays with size 2, and 1 subarray with size N. The time complexity is O(N^2).

A key observation is that the sum of a subarray [i, j] is equal to the sum of [0, j] minus the sum of [0, i - 1].

The sum of elements from index 0 to i is called the prefix sum (prefix = from the beginning, prefix sum = sum from the beginning, i.e. index 0). If we have the prefix sum for each index, we can find the sum of any subarray in constant time.

With the knowledge of the prefix sum under our belt, the problem boils down to Two Sum. We keep a dictionary of prefix_sum: index while going through the array calculating prefix_sum. If at any point, prefix_sum - target is in the dictionary, we have found our subarray.

Time Complexity: O(n)

Space Complexity: O(n)

First slide
First slide
First slide
First slide
First slide

Solution

1from typing import List
2
3def subarray_sum(arr: List[int], target: int) -> List[int]:
4    # prefix_sum 0 happens when we have an empty array
5    prefix_sums = {0: 0}
6    cur_sum = 0
7    for i in range(len(arr)):
8        cur_sum += arr[i]
9        complement = cur_sum - target
10        if complement in prefix_sums:
11            return [prefix_sums[complement], i + 1]
12        prefix_sums[cur_sum] = i + 1
13    return []
14
15if __name__ == "__main__":
16    arr = [int(x) for x in input().split()]
17    target = int(input())
18    res = subarray_sum(arr, target)
19    print(" ".join(map(str, res)))
20
1import java.util.Arrays;
2import java.util.HashMap;
3import java.util.List;
4import java.util.Scanner;
5import java.util.stream.Collectors;
6
7class Solution {
8    public static List<Integer> subarraySum(List<Integer> arr, int target) {
9        HashMap<Integer, Integer> prefixSums = new HashMap<>();
10        // prefix_sum 0 happens when we have an empty array
11        prefixSums.put(0, 0);
12        int curSum = 0;
13        for (int i = 0; i < arr.size(); i++) {
14            curSum += arr.get(i);
15            int complement = curSum - target;
16            if (prefixSums.containsKey(complement)) {
17                return List.of(prefixSums.get(complement), i + 1);
18            }
19            prefixSums.put(curSum, i + 1);
20        }
21        return null;
22    }
23
24    public static List<String> splitWords(String s) {
25        return s.isEmpty() ? List.of() : Arrays.asList(s.split(" "));
26    }
27
28    public static void main(String[] args) {
29        Scanner scanner = new Scanner(System.in);
30        List<Integer> arr = splitWords(scanner.nextLine()).stream().map(Integer::parseInt).collect(Collectors.toList());
31        int target = Integer.parseInt(scanner.nextLine());
32        scanner.close();
33        List<Integer> res = subarraySum(arr, target);
34        System.out.println(res.stream().map(String::valueOf).collect(Collectors.joining(" ")));
35    }
36}
37
1"use strict";
2
3function subarraySum(arr, target) {
4    // prefix_sum 0 happens when we have an empty array
5    const prefixSums = new Map([[0, 0]]);
6    let curSum = 0;
7    for (let i = 0; i < arr.length; i++) {
8        curSum += arr[i];
9        const complement = curSum - target;
10        if (prefixSums.has(complement)) {
11            return [prefixSums.get(complement), i + 1];
12        }
13        prefixSums.set(curSum, i + 1);
14    }
15    return [];
16}
17
18function splitWords(s) {
19    return s === "" ? [] : s.split(" ");
20}
21
22function* main() {
23    const arr = splitWords(yield).map((v) => parseInt(v));
24    const target = parseInt(yield);
25    const res = subarraySum(arr, target);
26    console.log(res.join(" "));
27}
28
29class EOFError extends Error {}
30{
31    const gen = main();
32    const next = (line) => gen.next(line).done && process.exit();
33    let buf = "";
34    next();
35    process.stdin.setEncoding("utf8");
36    process.stdin.on("data", (data) => {
37        const lines = (buf + data).split("\n");
38        buf = lines.pop();
39        lines.forEach(next);
40    });
41    process.stdin.on("end", () => {
42        buf && next(buf);
43        gen.throw(new EOFError());
44    });
45}
46
1#include <algorithm>
2#include <iostream>
3#include <iterator>
4#include <limits>
5#include <sstream>
6#include <string>
7#include <unordered_map>
8#include <vector>
9
10std::vector<int> subarray_sum(std::vector<int> arr, int target) {
11    std::unordered_map<int, int> prefix_sums;
12    // an empty subarray has a sum of 0
13    prefix_sums[0] = 0;
14    int cur_sum = 0;
15    for (int i = 0; i < arr.size(); i++) {
16        cur_sum += arr[i];
17        int complement = cur_sum - target;
18        if (prefix_sums.count(complement)) {
19            return {prefix_sums[complement], i + 1};
20        }
21        prefix_sums[cur_sum] = i + 1;
22    }
23    // return no indices if arr is empty
24    return {};
25}
26
27template<typename T>
28std::vector<T> get_words() {
29    std::string line;
30    std::getline(std::cin, line);
31    std::istringstream ss{line};
32    ss >> std::boolalpha;
33    std::vector<T> v;
34    std::copy(std::istream_iterator<T>{ss}, std::istream_iterator<T>{}, std::back_inserter(v));
35    return v;
36}
37
38void ignore_line() {
39    std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
40}
41
42template<typename T>
43void put_words(const std::vector<T>& v) {
44    if (!v.empty()) {
45        std::copy(v.begin(), std::prev(v.end()), std::ostream_iterator<T>{std::cout, " "});
46        std::cout << v.back();
47    }
48    std::cout << '\n';
49}
50
51int main() {
52    std::vector<int> arr = get_words<int>();
53    int target;
54    std::cin >> target;
55    ignore_line();
56    std::vector<int> res = subarray_sum(arr, target);
57    put_words(res);
58}
59

Note that in the implementation, typically we use prefix_sum[i] to denote the sum of elements in 0...i - 1 (the rightmost element i is not included in the sum). One good thing about this is prefix_sum[0] then means the sum of the array up to but not including the first element, i.e., empty array. The definition of an empty array sum is useful when there exists a subarray starting from 0 that sums up to the target. Without the definition of an empty array sum, we would miss it because its complement 0 does not exist in the dictionary.

Follow-up Problem

Find the total number of subarrays that sums up to target.

Explanation for the Follow-up Problem

The follow-up problem asks for the total number of subarrays that sum up to the target, rather than finding a single subarray. This slight change in the problem statement leads to a significant shift in our approach.

Key Differences:

  1. We're now counting all possible subarrays instead of finding just one.
  2. We need to keep track of frequency rather than indices.

New Approach:

Instead of storing "sum: index" in our hashmap, we now store "sum: frequency". Here's what this means:

  • Key (sum): The cumulative sum up to the current index.
  • Value (frequency): The number of times we've seen this particular sum.

Why this works:

  1. If at any point, our current cumulative sum minus the target (curr_sum - target) exists in the hashmap, it means we've found subarrays that sum to the target.
  2. The frequency stored for (curr_sum - target) tells us how many such subarrays end at the current index.

Example:

Consider the array [1, 2, 3] with a target of 3.

  1. Start with {0: 1} in the hashmap (sum of empty subarray).
  2. At index 0: sum = 1, add {1: 1} to hashmap.
  3. At index 1: sum = 3, it equals target, count = 1. Add {3: 1} to hashmap.
  4. At index 2: sum = 6, (6 - 3) = 3 exists in hashmap, count += 1. Update {6: 1} in hashmap.

Final count: 2 (subarrays [1, 2] and [3])

This approach allows us to count all subarrays efficiently in a single pass through the array, maintaining a time complexity of O(n) and space complexity of O(n).

1from typing import Counter, List
2
3def subarray_sum_total(arr: List[int], target: int) -> int:
4    prefix_sums: Counter[int] = Counter()
5    prefix_sums[0] = 1  # since empty array's sum is 0
6    cur_sum = 0
7    count = 0
8    for val in arr:
9        cur_sum += val
10        complement = cur_sum - target
11        if complement in prefix_sums:
12            count += prefix_sums[complement]
13        prefix_sums[cur_sum] += 1
14    return count
15
16if __name__ == "__main__":
17    arr = [int(x) for x in input().split()]
18    target = int(input())
19    res = subarray_sum_total(arr, target)
20    print(res)
21
1import java.util.Arrays;
2import java.util.HashMap;
3import java.util.List;
4import java.util.Scanner;
5import java.util.stream.Collectors;
6
7class Solution {
8    public static int subarraySumTotal(List<Integer> arr, int target) {
9        HashMap<Integer, Integer> prefixSums = new HashMap<>();
10        prefixSums.put(0, 1); // since empty array's sum is 0
11        int curSum = 0;
12        int count = 0;
13        for (int val : arr) {
14            curSum += val;
15            int complement = curSum - target;
16            if (prefixSums.containsKey(complement)) {
17                count += prefixSums.get(complement);
18            }
19            if (prefixSums.containsKey(curSum)) {
20                prefixSums.replace(curSum, prefixSums.get(curSum) + 1);
21            } else {
22                prefixSums.put(curSum, 1);
23            }
24        }
25        return count;
26    }
27
28    public static List<String> splitWords(String s) {
29        return s.isEmpty() ? List.of() : Arrays.asList(s.split(" "));
30    }
31
32    public static void main(String[] args) {
33        Scanner scanner = new Scanner(System.in);
34        List<Integer> arr = splitWords(scanner.nextLine()).stream().map(Integer::parseInt).collect(Collectors.toList());
35        int target = Integer.parseInt(scanner.nextLine());
36        scanner.close();
37        int res = subarraySumTotal(arr, target);
38        System.out.println(res);
39    }
40}
41
1"use strict";
2
3function subarraySumTotal(arr, target) {
4    const prefixSums = new Map();
5    prefixSums.set(0, 1); // since empty array's sum is 0
6    let curSum = 0;
7    let count = 0;
8    for (const val of arr) {
9        curSum += val;
10        const complement = curSum - target;
11        if (prefixSums.has(complement)) {
12            count += prefixSums.get(complement);
13        }
14        if (prefixSums.has(curSum)) {
15            prefixSums.set(curSum, prefixSums.get(curSum) + 1);
16        } else {
17            prefixSums.set(curSum, 1);
18        }
19    }
20    return count;
21}
22
23function splitWords(s) {
24    return s === "" ? [] : s.split(" ");
25}
26
27function* main() {
28    const arr = splitWords(yield).map((v) => parseInt(v));
29    const target = parseInt(yield);
30    const res = subarraySumTotal(arr, target);
31    console.log(res);
32}
33
34class EOFError extends Error {}
35{
36    const gen = main();
37    const next = (line) => gen.next(line).done && process.exit();
38    let buf = "";
39    next();
40    process.stdin.setEncoding("utf8");
41    process.stdin.on("data", (data) => {
42        const lines = (buf + data).split("\n");
43        buf = lines.pop();
44        lines.forEach(next);
45    });
46    process.stdin.on("end", () => {
47        buf && next(buf);
48        gen.throw(new EOFError());
49    });
50}
51
1#include <algorithm>
2#include <iostream>
3#include <iterator>
4#include <limits>
5#include <sstream>
6#include <string>
7#include <unordered_map>
8#include <vector>
9
10int subarray_sum_total(std::vector<int>& arr, int target) {
11    std::unordered_map<int, int> prefix_sums;
12    prefix_sums[0] = 1; // an empty array has a sum of 0
13    int cur_sum = 0;
14    int count = 0;
15    for (int val : arr) {
16        cur_sum += val;
17        int complement = cur_sum - target;
18        if (prefix_sums.count(complement)) {
19            count += prefix_sums[complement];
20        }
21        prefix_sums[cur_sum] += 1;
22    }
23    return count;
24}
25
26template<typename T>
27std::vector<T> get_words() {
28    std::string line;
29    std::getline(std::cin, line);
30    std::istringstream ss{line};
31    ss >> std::boolalpha;
32    std::vector<T> v;
33    std::copy(std::istream_iterator<T>{ss}, std::istream_iterator<T>{}, std::back_inserter(v));
34    return v;
35}
36
37void ignore_line() {
38    std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
39}
40
41int main() {
42    std::vector<int> arr = get_words<int>();
43    int target;
44    std::cin >> target;
45    ignore_line();
46    int res = subarray_sum_total(arr, target);
47    std::cout << res << '\n';
48}
49
Invest in Yourself
Your new job is waiting. 83% of people that complete the program get a job offer. Unlock unlimited access to all content and features.
Go Pro
Favorite (idle)