[Python | 2003] 수들의 합 2

https://www.acmicpc.net/problem/2003

---

문제 설명

N개의 수로 된 수열 A[1], A[2], …, A[N] 이 있다. 이 수열의 i번째 수부터 j번째 수까지의 합 A[i] + A[i+1] + … + A[j-1] + A[j]가 M이 되는 경우의 수를 구하는 프로그램을 작성하시오.

 

📌 시간 제한: 0.5초

---

입력

첫째 줄에 N(1 ≤ N ≤ 10,000), M(1 ≤ M ≤ 300,000,000)이 주어진다. 다음 줄에는 A[1], A[2], …, A[N]이 공백으로 분리되어 주어진다. 각각의 A[x]는 30,000을 넘지 않는 자연수이다.

---

출력

첫째 줄에 경우의 수를 출력한다.

---

문제 접근

이 문제는 수열에서 연속된 수들의 합이 특정 값 M이 되는 경우의 수를 찾는 전형적인 누적합 문제처럼 보일 수 있다. 하지만 모든 구간을 다 따져보면 시간복잡도가 O(N²)가 되기 때문에, N이 최대 10,000일 경우 효율적으로 해결할 수 없다.

그래서 이 문제에서는 "양의 정수"로만 이루어진 수열이라는 조건을 눈여겨봐야 한다. 이 조건 덕분에 투 포인터(Two Pointer) 기법을 사용할 수 있다. 투 포인터는 시작점과 끝점을 각각 가리키는 두 개의 인덱스를 두고, 이들이 수열을 스캔해 나가면서 조건을 만족하는 구간을 찾는 방식이다.

만약 음수가 포함된다면 포인터를 한쪽으로만 움직이면서 합을 조절할 수 없게 된다. 하지만 모든 수가 양수이기 때문에, 현재 합이 목표보다 작으면 구간을 확장하고(끝 포인터를 오른쪽으로), 현재 합이 목표 이상이면 구간을 줄이면서(시작 포인터를 오른쪽으로) 전체 합을 조절할 수 있다.

 

• 현재 구간의 합이 M보다 작다면, 오른쪽 포인터를 옮겨 구간을 확장한다.
• 현재 합이 M 이상이면, 왼쪽 포인터를 옮겨 구간을 줄이면서 조건을 맞춰본다.
• 합이 정확히 M일 때마다 카운트를 하나 증가시킨다.

 

또한 투 포인터는 각 포인터가 수열의 끝까지 최대 한 번씩만 이동하기 때문에 전체 시간복잡도는 O(N)이다. 이 덕분에 10,000개의 수가 들어와도 무난하게 처리할 수 있다.

---

구현

import sys
input = sys.stdin.readline

N, M = map(int, input().split())
arr = list(map(int, input().split()))

start, end = 0, 1
cnt = 0
current_sum = arr[0]

while end <= N:
    while current_sum >= M:
        if current_sum == M:
            cnt += 1
        current_sum -= arr[start]
        start += 1
    if end < N:
        current_sum += arr[end]
    end += 1

print(cnt)