[Leetcode] 560. Subarray Sum Equals K

문제

https://leetcode.com/problems/subarray-sum-equals-k/?envType=study-plan-v2&envId=top-100-liked

풀이

이 문제는 브루트포스로 접근하면 O(N^3)이 된다.

따라서 최적화를 위해 i~j구간의 합의 계산을 최적화하기 위해 누적합을 생각했고, 이를 통해 시간복잡도를 O(N^2)으로 개선했으니 시간초과에 걸렸다.

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        cum = [0,nums[0]]
        for num in nums[1:]:
            cum.append(cum[-1]+num)
        
        ans = 0
        for i in range(len(cum)):
            for j in range(i+1,len(cum)):
                if cum[j]-cum[i] == k: # i+1 ~ j까지의 구간합
                    ans += 1
                    
        return ans

이때, 생각해보니 누적합을 통해 i~j의 합을 구하여 k와 같은지 체크하는 공식은 누적합을 구한 배열을 cum이라 할 때 cum[j]-cum[i]=k인데, 이는 Two Sum 문제에서 접근했던 방식(hash table에 저장하여 linear한 탐색)과 동일한 접근법이 생각나서 이를 적용했다.

시간복잡도를 O(N)으로 개선해서 AC를 받을 수 있었다.

from collections import defaultdict

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        # 누적합 배열 계산해두기
        cum_list = [0,nums[0]]
        for num in nums[1:]:
            cum_list.append(cum_list[-1]+num)
        
        # 이후는 two sum 문제와 비슷하나 등장횟수를 세어줘야 한다.(경우의 수를 세야 하니)
        ans = 0
        cnt_dict = defaultdict(int) # 숫자: 등장횟수
        for cum_num in cum_list:
            if cum_num-k in cnt_dict:
                ans += cnt_dict[cum_num-k]
            cnt_dict[cum_num] += 1

        return ans

• 추가로 누적합을 통해 i~j의 합을 구할 때 배열에 제일 앞에 0을 넣어줘야 모든 구간에 대해 체크가 가능해진다.

• 구간합을 Prefix sum 또는 Cumulative sum이라고 호칭한다.
  

discussion을 보니 prefix sum을 먼저 쭉 구하고 체크하는게 아닌 prefix sum을 구하면서 체크할 수 있었다.

from collections import defaultdict

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        ans , prefsum= 0,  0
        d = defaultdict(int)
        d[0] = 1
        for num in nums:
            prefsum += num
            if  prefsum-k  in  d:
                ans = ans + d[prefsum-k]
            d[prefsum] += 1
        return ans