leetcode/subarray_equals_k.cpp at main · kevinsun203/leetcode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
#include <unordered_map>
#include <vector>

// Use array to keep track of all prefix sums. Iterate through the prefix array
// and find whether or not k can be obtained by prefixArray[i] - prefixArray[n]
// where n < i
class Solution {
public:

    int subarraySum(std::vector<int>& nums, int k) {
        std::unordered_map<int, int> prefixes;
        int length = nums.size();
        std::vector<int> prefixArray(length, nums[0]);
        int answer = 0;

        for (int i=1; i<length; i++) {
            prefixArray[i] = prefixArray[i-1] + nums[i];
        }

        for (int i=0; i<length; i++) {
            if (prefixArray[i] == k) {
                answer++;
            }

            int required = prefixArray[i] - k;
            if (prefixes.find(required) != prefixes.end()) {
                answer += prefixes[required];
            }

            if (prefixes.find(prefixArray[i]) != prefixes.end()) {
                prefixes[prefixArray[i]] += 1;
            } else {
                prefixes[prefixArray[i]] = 1;
            }
        }
        return answer;

    }
};