How to Count Subarrays with Sum K Using Prefix Sum and HashMap

Ctrip Employee Benefits

Ctrip, though not highly visible among internet giants, offers generous compensation and welfare, including a 12% housing fund, a "second-generation" child allowance of 10,000 CNY per child per year until age five, and three extra paid days per year for childcare.

Starting September 1, employees can work from home without needing direct supervisor approval; the request is automatically approved, reflecting a sincere effort to improve employee happiness.

LeetCode 560 – Subarray Sum Equals K

The problem asks for the number of subarrays whose sum equals a given integer k in an integer array nums.

Solution Overview

This classic problem can be solved using prefix sums combined with a hash map. By precomputing the prefix sum array sum, the number of subarrays ending at each index with sum k equals the count of previous prefix sums equal to sum[i] - k. The hash map records frequencies of prefix sums during a single pass.

Java Implementation

class Solution {
    public int subarraySum(int[] nums, int k) {
        int n = nums.length, ans = 0;
        int[] sum = new int[n + 10];
        for (int i = 1; i <= n; i++) sum[i] = sum[i - 1] + nums[i - 1];
        Map<Integer, Integer> map = new HashMap<>();
        map.put(0, 1);
        for (int i = 1; i <= n; i++) {
            int t = sum[i], d = t - k;
            ans += map.getOrDefault(d, 0);
            map.put(t, map.getOrDefault(t, 0) + 1);
        }
        return ans;
    }
}

C++ Implementation

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        int n = nums.size(), ans = 0;
        vector<int> sumv(n + 10, 0);
        for (int i = 1; i <= n; i++) sumv[i] = sumv[i - 1] + nums[i - 1];
        unordered_map<int, int> map;
        map[0] = 1;
        for (int i = 1; i <= n; i++) {
            int t = sumv[i], d = t - k;
            if (map.find(d) != map.end()) ans += map[d];
            map[t]++;
        }
        return ans;
    }
};

Python Implementation

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        n, ans = len(nums), 0
        sumv = [0] * (n + 10)
        for i in range(1, n + 1):
            sumv[i] = sumv[i - 1] + nums[i - 1]
        mapping = defaultdict(int)
        mapping[0] = 1
        for i in range(1, n + 1):
            t = sumv[i]
            d = t - k
            ans += mapping[d]
            mapping[t] += 1
        return ans

TypeScript Implementation

function subarraySum(nums: number[], k: number): number {
    let n = nums.length, ans = 0;
    const sum = new Array<number>(n + 10).fill(0);
    for (let i = 1; i <= n; i++) sum[i] = sum[i - 1] + nums[i - 1];
    const map = new Map<number, number>();
    map.set(0, 1);
    for (let i = 1; i <= n; i++) {
        const t = sum[i], d = t - k;
        if (map.has(d)) ans += map.get(d)!;
        map.set(t, (map.get(t) ?? 0) + 1);
    }
    return ans;
};

Complexity Analysis

Time complexity: O(n) for computing prefix sums and scanning the array.

Space complexity: O(n) for the prefix sum array and hash map.