파이썬으로 이진 리스트에서 합이 k인 부분 리스트의 개수 구하기

0과 1로만 이루어진 이진 리스트가 있다고 가정해 봅시다. 여기에 또 다른 입력값 k가 주어졌을 때, 합이 k와 같은 연속된 부분 리스트의 개수를 구하는 것이 목표입니다.

예를 들어, 입력이 nums = [1, 0, 0, 1, 1, 1, 0, 1], k = 3이라면 출력은 8이 됩니다. 조건을 만족하는 부분 리스트가 [1,0,0,1,1], [0,0,1,1,1], [0,0,1,1,1,0], [0,1,1,1], [0,1,1,1,0], [1,1,1], [1,1,1,0], [1,1,0,1]로 총 8개이기 때문입니다.

해결 방법

이 문제는 누적 합(prefix sum)과 해시 맵을 활용하면 효율적으로 풀 수 있습니다. 절차는 다음과 같습니다.

sums := 키 0에 값 1을 담고 있는 맵으로 초기화
r_sum := 0 (현재까지의 누적 합)
ans := 0 (정답 카운트)
nums의 각 요소 x에 대해 반복:
- r_sum := r_sum + x
- ans := ans + (r_sum - k가 sums에 존재하면 sums[r_sum - k], 없으면 0)
- sums[r_sum] := sums.get(r_sum, 0) + 1 (현재 누적 합의 등장 횟수 증가)
ans 반환

동작 원리

핵심 아이디어는 다음과 같습니다. 현재 위치까지의 누적 합을 r_sum이라 할 때, 이전 어느 시점까지의 누적 합이 r_sum - k였다면, 그 시점 다음부터 현재 위치까지의 부분 리스트의 합은 정확히 k가 됩니다. 따라서 맵 sums에 각 누적 합이 등장한 횟수를 기록해 두면, 현재 위치에서 조건을 만족하는 부분 리스트의 개수를 상수 시간에 바로 알 수 있습니다. 초기값으로 sums[0] = 1을 설정하는 이유는 리스트의 처음부터 시작하는 부분 리스트도 함께 고려하기 위함입니다. 이 알고리즘의 시간 복잡도는 O(n)입니다.

예제 코드

다음 구현을 통해 더 자세히 이해해 보겠습니다.

def solve(nums, k):
    sums = {0: 1}
    r_sum = 0
    ans = 0

    for x in nums:
        r_sum += x
        ans += sums.get(r_sum - k, 0)
        sums[r_sum] = sums.get(r_sum, 0) + 1

    return ans

nums = [1, 0, 0, 1, 1, 1, 0, 1]
k = 3
print(solve(nums, k))

입력

[1, 0, 0, 1, 1, 1, 0, 1], 3