파이썬으로 합이 k로 나누어 떨어지는 연속 부분 수열 개수 구하는 방법

문제 이해하기

배열 nums와 값 k가 주어졌을 때, 합이 k로 나누어 떨어지는 연속 부분 수열(연속된 원소로 이루어진 부분 배열)의 개수를 구해야 합니다.

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

접근 방법: 누적 합의 나머지 활용

모든 부분 배열을 일일이 확인하는 브루트 포스 방식은 O(n²)의 시간이 걸리지만, 누적 합(prefix sum)의 나머지를 활용하면 O(n)으로 효율적으로 해결할 수 있습니다.

핵심 아이디어는 다음과 같습니다. 두 지점의 누적 합을 k로 나눈 나머지가 서로 같다면, 그 두 지점 사이에 있는 부분 배열의 합은 반드시 k로 나누어 떨어집니다.

풀이 단계는 다음과 같습니다.

• 크기가 k인 배열 x를 만들고 모든 값을 0으로 초기화합니다. x[i]는 '지금까지 등장한 누적 합 중 k로 나눈 나머지가 i인 경우의 수'를 저장합니다.
• x[0] := 1로 설정합니다. 빈 접두사를 고려하여, 배열의 처음부터 시작하는 부분 배열도 올바르게 세기 위함입니다.
• r := 0, s := 0으로 초기화합니다. r은 정답 개수, s는 현재까지 누적 합을 k로 나눈 나머지입니다.
• nums의 각 원소에 대해 다음을 반복합니다.
  • s := (s + elem) mod k — 누적 합의 나머지를 갱신합니다.
  • r := r + x[s] — 같은 나머지를 가진 이전 지점의 수만큼 정답에 더합니다.
  • x[s] := x[s] + 1 — 현재 나머지의 등장 횟수를 1 증가시킵니다.
• r을 반환합니다.

구현 예제

아래 파이썬 코드로 직접 확인해 보겠습니다.

def solve(k, nums):
   x = [0]*k
   x[0] = 1
   r = s = 0
   for elem in nums:
      s = (s + elem) % k
      r += x[s]
      x[s] += 1
   return r

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

입력

k = 3, nums = [1,2,3,4,1]

출력

4

동작 과정 살펴보기

예제 입력 [1,2,3,4,1]에서 알고리즘이 어떻게 진행되는지 단계별로 살펴보겠습니다.

최종 결과는 r = 4로, 앞서 찾은 네 가지 부분 수열과 정확히 일치합니다.

복잡도 분석

시간 복잡도: O(n) — 배열을 한 번만 순회하면 됩니다.
공간 복잡도: O(k) — 크기 k의 나머지 카운트 배열 하나만 필요합니다.

이처럼 누적 합과 나머지의 성질을 활용하면, 겹치는 계산 없이 문제를 선형 시간 안에 해결할 수 있습니다. 유사한 패턴은 '합이 k인 부분 배열 찾기' 등 다양한 누적 합 문제에서도 응용되므로 꼭 기억해 두시길 바랍니다.