560. Subarray Sum Equals K (Array, Hash Table, Prefix Sum)

링크 - https://leetcode.com/problems/subarray-sum-equals-k/description/?envType=company&envId=facebook&favoriteSlug=facebook-thirty-days

난이도 - Medium

Intuition

이 문제는 처음에는 2중 for문으로 접근하는게 먼저 떠오를 것이다. (3중 for문을 쓰는 Brute Force는 여기서 일단 배제하자). 2중 for문을 사용하면 매 for문마다 시작하는 부분부터 나오는 값들을 총 합해 그 합이 우리가 찾고자 하는 값과 일치하는지를 확인하는 것이다.

하지만 이 방법을 사용하면 Time Exceed가 나오기 때문에 최대한 더 효율적인 Time Complexity를 갖는 알고리즘이 뭔지 생각해봐야 할 것이다.

 

따라서 우리는 여기서 Hash Table를 사용해 O(n)인 Time complexity를 얻을 수 있다. 기본적인 아이디어는 i부터 j까지의 부분집합의 총 합을 알고 싶다면, 0부터 j까지의 부분집합의 합에서 0부터 i-1까지의 부분집합의 합을 빼면 된다는 것에서부터 시작한다.

1. hashmap을 초기화 하고, key-value가 0,1인 pair를 집어넣는다. 이는 나중에 원소 하나의 값이 우리가 찾는 k와 같을때를 대비해서 넣는 것이다.

2. 첫번째 원소부터 for문을 통해 탐색한다. 원소을 탐색할때마다 총 합을 계산하는 변수에 현재 원소값을 더한다. 따라서 cur_sum += nums[i]이다. 

3. 그 다음, 해시테이블에 cur_sum - k값의 key를 가지고 있는지 확인한다. 그 이유는 우리가 현재까지 cur_sum을 통해 알고있는 0부터 i까지의 부분집합의 총합에서, 0부터 j까지 cur_sum - k만큼의 부분집합의 총합을 가지고 있는것을 빼면, 우리는 j+1부터 i까지의 합은 k임을 알 수 있는것이기 때문이다. 따라서 ans += hashmap.get(cur_sum - k)를 한다.

4. 그 다음, 해시테이블에 0부터 현재 인덱스까지의 총 합을 집어넣는다. 그래야 나중에 이 값을 비교하고 ans를 업데이트 할 수 있기 때문이다. 물론 만약에 이미 해시테이블에 같은 key가 있다면, 1를 더해서 넣어주는것을 잊지말자. 따라서 hashmap[cur_sum] = hahsmap.get(cur_sum, 0) + 1

5. 2-4를 반복하고 for문이 끝나면 ans를 반환한다.

Code

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        ans = 0
        cur_sum = 0

        hashmap = {}
        hashmap[0] = 1

        for i in range(len(nums)):
            cur_sum += nums[i]
            if cur_sum - k in hashmap:
                ans += hashmap.get(cur_sum - k)
            hashmap[cur_sum] = hashmap.get(cur_sum, 0) + 1

        return ans

Complexity

Time Complexity: O(n) For문을 통해 nums를 돌기 때문에 O(n)이다. 

Space Complexity: O(n) Hashmap은 최대 n개의 원소를 가지고 있을 수 있다.