[백준] 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을 넘지 않는 자연수이다.

출력

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

풀이

Two Pointer를 사용!

두 개의 포인터를 이용해서 탐색하는 것으로 "특정한 합을 가지는 부분 연속 수열" 문제에서 사용한다.

 

답안 코드

예외 처리를 생각하지 못해서 조금 시간이 걸렸지만... 해결해냈다!

import sys
input = sys.stdin.readline

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

left = 0
right = 1

answer = 0

result = arr[0]

while True:
    if result < M and right == N: break

    if result < M and right < N:
        result+=arr[right]
        right+=1
    else:
        while left <= right and result >= M:
            if result == M:
                answer+=1
            result-=arr[left]
            left+=1
print(answer)

결과

다른 사람들의 코드를 보니,

result로 값을 빼고 더하지 않고 sum(arr[right:left])을 사용하는 것 같았다.

이렇게 하면 조금 더 직관적이고 간결한 코드가 작성되었다. 

 

import sys
input = sys.stdin.readline

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

start = 0
end = 0

answer = 0

while end <= N:
    result = sum(arr[start:end])
    
    if result < M:
        end+=1
    elif result > M:
        start+=1

    if result == M:
        answer+=1
        start+=1
        end+=1
    
print(answer)

결과

하지만 역시 매번 배열의 합을 구하다 보니 시간이 7~8배 더 느리다는 점... 

 

결론

"특정한 합을 가지는 부분 연속 수열" 문제에서는
Two Pointer를 사용하자!