[백준/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이 되는 경우의 수를 구하는 프로그램을 작성하시오.

입력

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

출력

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

코드

def count_subarrays_with_sum(N, M, A):
    start, end = 0, 0  
    current_sum = 0  
    count = 0 

    while end < N:
        current_sum += A[end]
        end += 1

        while current_sum > M and start < end:
            current_sum -= A[start]
            start += 1

        if current_sum == M:
            count += 1

    return count

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

print(count_subarrays_with_sum(N, M, A))

코드 풀이

start와 end를 정의하고 투 포인터로 이용
current_sum: 현재 부분합, count : 경우의 수로 정의
부분합이 M보다 작은 경우는 end를 오른쪽으로 이동하며 합을 증가, 
부분합이 M 이상인 경우는 start를 오른쪽으로 이동하며 합을 감소,
현재 부분합이 M인 경우는 경우의 수 증가함
count 를 반환하여 전체의 경우의 수를 계산함