[Python/파이썬] 백준 2003번 수들의 합 2

2003번: 수들의 합 2

 

문제

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

입력

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

출력

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

코드

n, m = map(int, input().split())
arr = list(map(int, input().split()))

s, e = 0, 0
partial_sum = arr[0]
ans = 0
while e < n:
    if partial_sum < m:
        e += 1
        if e >= n:
            break
        partial_sum += arr[e]
    else:
        if partial_sum == m:
            ans += 1
        partial_sum -= arr[s]
        s += 1
print(ans)

합을 계산하고자 하는 부분배열의 시작점과 끝점을 가리키는 포인터 s, e를 만들고, 두 포인터의 위치를 조절해가며 부분배열의 합이 m이 되는 경우의 수들을 구한다.

만약 현재 arr[s]~arr[e]까지의 합(partial_sum)이 m보다 작다면 더 많은 원소가 필요하므로 e를 한 칸 뒤로 보낸다.

그것이 아니고 partial_sum이 m보다 작거나 같다면 s를 한 칸 뒤로 보내서 포함되는 원소가 1개 줄어들도록 한다. 이때, partial_sum이 m과 같은 경우에는 ans에 1개를 더해준다.