[Leetcode] Subarray Sum Equals K

Problem)

 

Solution)

쉬워 보이지만 제대로된 풀이법을 생각해 내는데 꾀 오랜 시간이 걸렸다..

 

계속해서 무식한 방법을 사용하다 보니까 시간 초과가 나서...

 

leetcode 사이트에 discussion에 좋은 설명이 있어서 참조하여 설명글을 정리하겠습니다.

 

기본 아이디어는 입력받은 리스트 값을 하나씩 단계별로 더해서 sum값을 구할 때, 이 sum값이 k만큼 증가하는 순간에 합이 k값인 subarray를 구할 수 있습니다. 무슨 의미인지 풀어서 설명하겠습니다.

 

리스트 [1, 2, 1, 3]과 k = 3이 주어진다면, 단계별 sum값들은 [1, 3, 4, 7]이 됩니다. 이때, 1 -> 4인 순간과 4 -> 7인 순간을 통해서 우리는 합이 3인 subarray가 2개 있다는 것을 알 수 있습니다.

 

하지만, 실제 리스트 [1, 2, 1, 3]에서 합이 3인 subarray는 [1,2], [2,1], [3]으로 총 3개가 됩니다. 위의 계산에서 [3]을 빠뜨렸다는 것을 알 수 있습니다. 따라서 단계별 합 값이 k가 되는(위의 예제에서는 [3]) 경우를 포함하기 위해서 우리는 합계에 0을 추가하여 고려해야 합니다.

 

따라서, [0, 1, 3, 4, 7]을 순서대로 계산했을 때, 0 -> 3, 1 -> 4, 4 -> 7 총 3개의 subarray가 있다는 것을 알 수 있습니다.

 

import collections

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        # 리스트 값의 합 카운트 사전
        sums_list = collections.Counter()
        # k와 같은 합 값이 존재할 경우를 위해
        sums_list[0] = 1
        
        # 리스트 값의 합
        sums = 0
        
        # 결과값
        answer = 0
        
        for num in nums:
            # 리스트 값을 하나씩 더하며 이때까지의 합 값 구하기
            sums += num
            
            # 만약 이전에 sums값 중 (sums-k)와 같은 값이 존재 했다면,
            # 합이 k인 subarray가 존재한다는 의미
            answer += sums_list[sums-k]
            
            sums_list[sums] += 1

        return answer