930. Binary Subarrays With Sum

[zeikar]

Problem link

https://leetcode.com/problems/binary-subarrays-with-sum/

Problem Summary

부분 배열 중 합이 goal인 개수를 출력하는 문제

Solution

뭔가 슬라이딩 윈도우로 되지 않을까 했는데 쉽지 않다.
가능하긴 한데 0 prefix 개수를 관리하거나 (1 pass) 최대 k인 개수를 구한 다음 최대 k-1인 개수를 빼서 구하거나 (2 pass)... 에디토리얼 참고.

prefix sum을 쓰면 쉽게 풀린다.
지금까지의 합이 나온 수를 카운트 하면서 현재까지의 합 - goal 의 개수를 계속해서 더해주면 된다.
현재까지 goal을 만들 수 있는 부분 배열의 수는 sum-goal을 만들 수 있는 경우의 수를 더한다고 생각할 수 있다.

Source Code

class Solution:
    def numSubarraysWithSum(self, nums: List[int], goal: int) -> int:
        ans = 0
        sum = 0
        cnt = defaultdict(int)
        cnt[0] = 1

        for num in nums:
            sum += num
            ans += cnt[sum - goal]
            cnt[sum] += 1

        return ans