백준 1806번 : 부분합 in Python

1806번: 부분합

 

www.acmicpc.net

이 문제는 자연수로 이루어진 수열에서 연속된 수들의 부분합이 특정 수를 넘는 것 중 가장 짧은 것의 길이를 구하는 문제이다. 

가장 단순한 방법으로는 모든 가능한 연속된 부분수열에 대해 합을 구하고 이 합이 특정 수를 넘으면 그 부분수열의 길이를 구해 그 중 가장 짧은 것을 찾는 방법이 있다. 하지만 이 방법은 길이가 N인 수열의 모든 연속된 수로 이루어진 부분수열의 수가 N^2개이고, 각 부분수열의 합을 구하는데 O(N)만큼의 시간이 소요되기에 총 O(N^3)의 시간이 소요된다는 단점이 있다.

이를 효율적으로 만들기 위해 앞에서 사용했던 투 포인터를 이용한다. 여기서 두 개의 포인터는 연속된 수로 이루어진 부분수열의 양 끝을 의미한다. 이 문제에서는 우선 두 개의 포인터(l, r)를 모두 왼쪽 끝에 위치시킨다. 이로 인해 만들어진 부분수열의 합을 구한 뒤 만약 합이 기준이 되는 수보다 작은 경우 r 포인터를 오른쪽으로 한 칸씩 이동시키고 이에 따라 합을 수정한다. 반대로 합이 기준이 되는 수보다 작은 경우 그 때의 부분수열의 길이를 구해 조건을 만족하는 가장 짧은 부분수열의 길이보다 더 작은 경우 이를 수정한다. 그리고 l 포인터를 오른쪽으로 한 칸씩 이동시키고 이에 따라 합을 수정한다.

이때 매번 포인터를 한 칸씩 옮길 때마다 부분수열의 합을 구할 필요없이 이전에 구해둔 부분수열의 합에 추가된 원소를 더하거나 삭제된 원소를 빼면 부분합을 구하는데 걸리는 시간을 O(1)으로 줄일 수 있다. 이를 r이 N-1보다 작거나 같고, l이 r보다 작거나 같은 경우 동안 계속 반복하면 된다. 단, r이 N-1이 됐을 때 부분합이 기준의 되는 수보다 작으면 부분수열에 새로운 수를 추가하는 것이 불가능하므로 반복을 중단한다.

이렇게 해서 얻은 부분합이 만약 존재한다면 이를 출력하고, 만약 부분합을 구하지 못한 경우에는 0을 출력한다.

이를 코드로 구현하면 다음과 같다. 

import sys
input = sys.stdin.readline
N, S = map(int, input().split())
arr = list(map(int, input().split()))
# 두 개의 포인터를 왼쪽 끝에 위치시킨다.
l = 0
r = 0
# 두 개의 포인터 사이의 모든 수를 더한 부분합을 저장할 변수
s = arr[0]
# 부분수열 중 가장 짧은 것의 길이를 저장할 변수
min_len = sys.maxsize
while r <= N-1 and l <= r:
    # 부분합이 기준이 되는 수보다 크거나 같은 경우
    if s >= S:
        min_len = min(min_len, r-l+1)
        s -= arr[l]
        l += 1
    else:
        if r < N-1:
            r += 1
            s += arr[r]
        # 부분합이 기준이 되는 수보다 작으면서 r이 N-1이면 
        # 더 이상 부분합을 증가시킬 수 없으므로 while loop을 빠져나온다.
        else:
            break
# 어떤 부분합도 기준이 되는 수를 넘지 못하는 경우
if min_len == sys.maxsize: print(0)
else: print(min_len)