[leetcode #560] Subarray Sum Equals K

Problem

Given an array of integers nums and an integer k, return the total number of continuous subarrays whose sum equals to k.

Example 1:

Input: nums = [1,1,1], k = 2
Output: 2

Example 2:

Input: nums = [1,2,3], k = 3
Output: 2

Constraints:

・ 1 <= nums.length <= 2 * 10⁴
・ -1000 <= nums[i] <= 1000
・ -10⁷ <= k <= 10⁷

Idea

2021년 4월에 출제됐던 문제가 다시 등장했다.

주어진 배열에서 subarray의 합이 k와 일치하는 경우가 몇 개가 되는지 구하는 문제다.

Brute-force로 매번 subarray의 합을 구하면 O(n³)이 되므로 Time Limit Exceeded가 걸린다. 이를 해결하기 위한 방법으로 배열을 탐색하면서 i번째 index에서 현재 index까지의 합을 저장하는 배열 (sums)을 사용할 수 있다. 현재 index에서 배열의 값을 sums의 각 index에 더하고 더한 값이 k와 일치할 경우 count를 1씩 올린다. 배열의 탐색이 끝난 뒤 count를 리턴하면 된다.

Time Complexity: O(n²)
Space Complexity: O(n)

이보다 더 좋은 방법은 HashMap을 이용하는 것이다. running sum을 S라고 할 때 S(i) - S(j) = k 를 만족하는 index i와 j가 있다면 합이 k인 subarray가 존재하기 때문에 HashMap을 사용할 수 있다. 배열을 탐색하면서 running sum을 구하고 running sum을 HashMap에 더한다. HashMap에 해당 값이 있다면 count를 1씩 올리고 없다면 새로운 key로 만들어 값을 1로 해서 HashMap에 더한다. Running sum을 구하면서 현재 running sum S(i)에서 k를 뺀 값이 HashMap에 존재한다면 합이 k인 subarray가 존재한다는 뜻이다. 이 때 map에서 count 값을 얻어와 결과값에 더하면 된다.

Time Complexity: O(n)
Space Complexity: O(n)

Solution

dp

class Solution {
    public int subarraySum(int[] nums, int k) {
        int[] sums = new int[nums.length];
        int res = 0;

        for (int i=0; i < nums.length; i++) {
            for (int j=0; j <= i; j++) {
                sums[j] += nums[i];
                if (sums[j] == k) {
                    res++;
                }
            }
        }

        return res;
    }
}

HashMap

class Solution {
    public int subarraySum(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();
        map.put(0, 1);

        int sum = 0;
        int res = 0;

        for (int i=0; i < nums.length; i++) {
            sum += nums[i];
            if (map.containsKey(sum - k)) {
                res += map.get(sum - k);
            }
            map.put(sum, map.getOrDefault(sum, 0) + 1);
        }

        return res;
    }
}

Reference

https://leetcode.com/problems/subarray-sum-equals-k/