[Python/파이썬] 백준 1806번 부분합

1806번: 부분합

문제

10,000 이하의 자연수로 이루어진 길이 N짜리 수열이 주어진다. 이 수열에서 연속된 수들의 부분합 중에 그 합이 S 이상이 되는 것 중, 가장 짧은 것의 길이를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 N (10 ≤ N < 100,000)과 S (0 < S ≤ 100,000,000)가 주어진다. 둘째 줄에는 수열이 주어진다. 수열의 각 원소는 공백으로 구분되어져 있으며, 10,000이하의 자연수이다.

출력

첫째 줄에 구하고자 하는 최소의 길이를 출력한다. 만일 그러한 합을 만드는 것이 불가능하다면 0을 출력하면 된다.

코드

import sys
input = sys.stdin.readline

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

acc = [0]
# 누적합 구하기
for i in range(1, n + 1):
    acc.append(acc[i-1] + arr[i-1])

start, end = 0, 0
ans = 100001
while start < n: 
    if acc[end] - acc[start] < s:
        if end == n:
            break
        end += 1
    else:
        ans = min(ans, end - start)
        start += 1

if ans == 100001:        
    print(0)
else:
    print(ans)

보고 누적합과 투포인터가 생각났던 문제였다.

이렇게 부분합을 구하는 문제를 풀이할 때는 매번 더하기보다는 누적합으로 한 번에 구해놓고 하는 것이 더 효율적이라고 이전에 다른 문제에서 봤던 기억이 있다. 문제의 풀이 방법은 간단하다.

1. 입력받은 배열의 누적합을 우선 구해둔다. 이렇게 하면 연속된 수들의 부분합을 뺄셈 한 번으로 편하게 구할 수 있다.
2. 부분합이 s보다 작다면 end를 증가시키고, s보다 크다면 그때의 길이를 이미 저장되어있던 최소값과 비교하여 더 작은 것을 저장하는 과정을 반복한다. 이때 부분합이 s보다 작은데 end가 이미 끝에 도달했다면 start를 계속 증가시켜봤자 부분합은 감소하기만 하므로 더이상 볼 필요가 없으니 break한다.

이렇게 풀이해놓고도 계속 틀려서 원인을 찾아봤더니

10 10

1 1 1 1 1 1 1 1 1

이 케이스를 고려해주지 않아서 계속 틀렸었다. 원래 풀었던 코드에서는 모두 다 더해야 s 이상을 만들 수 있는 경우를 고려하지 못해 이 경우는 만들 수 없는 것으로 인식했었다. 이를 해결하기 위해 누적합을 저장하는 배열인 acc의 0번째 값으로는 0을 넣어주었다. 입력으로 받은 배열 arr의 누적합을 1번째 값부터 채워서 이러한 경우도 처리할 수 있도록 하였다.