[PS] BOJ 1806 / 부분합

[PS] BOJ 1806 / 부분합
문제 링크: https://www.acmicpc.net/problem/1806
Thumbnail: Photo by Thibault Penin (Unsplash)

전체 수열의 부분 수열 중, 전체 원소의 합이 $S$보다 크거나 같으면서 길이가 최소가 되는 수열을 찾아, 그 길이를 구하는 문제입니다.

풀이

0.5초라는 짧은 시간 제한 때문에, 모든 부분 수열을 탐색할 수 없습니다. 따라서, 투 포인터 방식으로 유망한 경우만 빠르게 탐색해야 합니다.

누적 합 계산

전체 배열을 입력받으면서 동시에 누적 합 배열도 계산했습니다. 누적 합 계산의 편의를 위해, 입력 배열과 누적합 배열의 0번 칸에는 0을 채워넣고 1번 칸부터 원소를 채웠습니다.

N, S = map(int, input().split())
arr = [0]
prefix_sum = [0]
for i, e in enumerate(map(int, input().split()), 1):
    arr.append(e)
    prefix_sum.append(prefix_sum[i - 1] + e)

누적합 계산

투 포인터

기본적으로, (1, 1)에서 시작해 구간을 옮겨가며 탐색합니다.

구간의 시작 위치를 l, 끝 위치를 r이라고 하면

  • r이 $N$보다 작거나 같을 동안 반복하며
    • 현재 구간의 길이와 누적 합을 계산합니다.
    • 현재 구간의 합이 $S$보다 크거나 같다면,
      • 만약 현재 구간의 길이가 기존에 찾아둔 최소 길이보다 짧다면 갱신합니다.
      • 이후, l을 1 증가시킵니다. (더 짧은 경우를 찾기 위함)
    • 현재 구간의 합이 $S$보다 작다면, r을 1 증가시킵니다. (구간 합을 키워야 함)
total = 0
min_len = N + 1
l = 1
r = 1
while r <= N:
    total = prefix_sum[r] - prefix_sum[l - 1]
    length = r - l + 1
    if total >= S:
        if length < min_len:
            min_len = length
        l += 1
    else:
        r += 1

print(0 if min_len == N + 1 else min_len)

투 포인터 탐색

전체 코드

input = open(0).readline
N, S = map(int, input().split())
arr = [0]
prefix_sum = [0]
for i, e in enumerate(map(int, input().split()), 1):
    arr.append(e)
    prefix_sum.append(prefix_sum[i - 1] + e)

total = 0
min_len = N + 1
l = 1
r = 1
while r <= N:
    total = prefix_sum[r] - prefix_sum[l - 1]
    length = r - l + 1
    if total >= S:
        if length < min_len:
            min_len = length
        l += 1
    else:
        r += 1

print(0 if min_len == N + 1 else min_len)

solution.py