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. Time complexity is O(N^2).

A key observation is that the 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. If we have the subarray 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
14if __name__ == '__main__':
15    arr = [int(x) for x in input().split()]
16    target = int(input())
17    res = subarray_sum(arr, target)
18    print(' '.join(map(str, res)))
19
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
1function subarraySum(arr, target) {
2    // prefix_sum 0 happens when we have an empty array
3    const prefixSums = new Map([[0, 0]]);
4    let curSum = 0;
5    for (let i = 0; i < arr.length; i++) {
6        curSum += arr[i];
7        const complement = curSum - target;
8        if (prefixSums.has(complement)) {
9            return [prefixSums.get(complement), i + 1];
10        }
11        prefixSums.set(curSum, i + 1);
12    }
13}
14
15function splitWords(s) {
16    return s == "" ? [] : s.split(' ');
17}
18
19function* main() {
20    const arr = splitWords(yield).map((v) => parseInt(v));
21    const target = parseInt(yield);
22    const res = subarraySum(arr, target);
23    console.log(res.join(' '));
24}
25
26class EOFError extends Error {}
27{
28    const gen = main();
29    const next = (line) => gen.next(line).done && process.exit();
30    let buf = '';
31    next();
32    process.stdin.setEncoding('utf8');
33    process.stdin.on('data', (data) => {
34        const lines = (buf + data).split('\n');
35        buf = lines.pop();
36        lines.forEach(next);
37    });
38    process.stdin.on('end', () => {
39        buf && next(buf);
40        gen.throw(new EOFError());
41    });
42}
43
1#include <algorithm> // copy
2#include <iostream> // cin, cout, streamsize
3#include <iterator> // back_inserter, istream_iterator, ostream_iterator, prev
4#include <limits> // numeric_limits
5#include <sstream> // istringstream
6#include <string> // getline, string
7#include <unordered_map> // unordered_map
8#include <vector> // 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    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
41template<typename T>
42void put_words(const std::vector<T>& v) {
43    if (!v.empty()) {
44        std::copy(v.begin(), std::prev(v.end()), std::ostream_iterator<T>{std::cout, " "});
45        std::cout << v.back();
46    }
47    std::cout << '\n';
48}
49
50int main() {
51    std::vector<int> arr = get_words<int>();
52    int target;
53    std::cin >> target;
54    ignore_line();
55    std::vector<int> res = subarray_sum(arr, target);
56    put_words(res);
57}
58

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

Follow up

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

Explanation

Since the new problem does not ask for index but total number instead, we can change our hashmap to "sum k: number of prefix sums that sums up to k".

1from typing import Counter, List
2
3def subarray_sum_total(arr: List[int], target: int) -> int:
4    prefix_sums = Counter()
5    prefix_sums[0] = 1 # since empty array's sum is 0
6    cur_sum = 0
7    count = 0
8    for i in range(len(arr)):
9        cur_sum += arr[i]
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 i = 0; i < arr.size(); i++) {
14            curSum += arr.get(i);
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
1function subarraySumTotal(arr, target) {
2    const prefixSums = new Map();
3    prefixSums.set(0, 1);  // since empty array's sum is 0
4    let curSum = 0;
5    let count = 0;
6    for (let i = 0; i < arr.length; i++) {
7        curSum += arr[i];
8        const complement = curSum - target;
9        if (prefixSums.has(complement)) {
10            count += prefixSums.get(complement);
11        }
12        if (prefixSums.has(curSum)) {
13            prefixSums.set(curSum, prefixSums.get(curSum) + 1);
14        } else {
15            prefixSums.set(curSum, 1);
16        }
17    }
18    return count;
19}
20
21function splitWords(s) {
22    return s == "" ? [] : s.split(' ');
23}
24
25function* main() {
26    const arr = splitWords(yield).map((v) => parseInt(v));
27    const target = parseInt(yield);
28    const res = subarraySumTotal(arr, target);
29    console.log(res);
30}
31
32class EOFError extends Error {}
33{
34    const gen = main();
35    const next = (line) => gen.next(line).done && process.exit();
36    let buf = '';
37    next();
38    process.stdin.setEncoding('utf8');
39    process.stdin.on('data', (data) => {
40        const lines = (buf + data).split('\n');
41        buf = lines.pop();
42        lines.forEach(next);
43    });
44    process.stdin.on('end', () => {
45        buf && next(buf);
46        gen.throw(new EOFError());
47    });
48}
49
1#include <algorithm> // copy
2#include <iostream> // cin, cout, streamsize
3#include <iterator> // back_inserter, istream_iterator
4#include <limits> // numeric_limits
5#include <sstream> // istringstream
6#include <string> // getline, string
7#include <unordered_map> // unordered_map
8#include <vector> // 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, count = 0;
14    for (int i = 0; i < arr.size(); i++) {
15        cur_sum += arr[i];
16        int complement = cur_sum - target;
17        if (prefix_sums.count(complement)) {
18            count += prefix_sums[complement];
19        }
20        if (prefix_sums.count(cur_sum)) {
21            prefix_sums[cur_sum] += 1;
22        } else {
23            prefix_sums[cur_sum] = 1;
24        }
25    }
26    return count;
27}
28
29template<typename T>
30std::vector<T> get_words() {
31    std::string line;
32    std::getline(std::cin, line);
33    std::istringstream ss{line};
34    std::vector<T> v;
35    std::copy(std::istream_iterator<T>{ss}, std::istream_iterator<T>{}, std::back_inserter(v));
36    return v;
37}
38
39void ignore_line() {
40    std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
41}
42
43int main() {
44    std::vector<int> arr = get_words<int>();
45    int target;
46    std::cin >> target;
47    ignore_line();
48    int res = subarray_sum_total(arr, target);
49    std::cout << res << '\n';
50}
51

Still not clear? Submit the part you don't understand to our editors.