1806 - 부분합

이 문제는 자연수 N개로 이루어진 수열이 주어졌을 때, S 이상을 만족하는 수열의 부분합 중 최소 길이를 찾아야 한다.

문제를 어떻게 풀 수 있을까? 필자의 경우 가장 먼저 떠올린 방법은 누적합과 BST를 이용하는 방법이었다. 첫 원소부터 누적합을 구하면서, 구한 누적합을 트리에 넣는다. 그런 다음 각 누적합 원소마다 S 이상인지 확인한다. 만약 S 이상이라면 누적합 원소 - S보다 작거나 같은 누적합 원소 중 가장 큰 원소를 찾아 그 원소의 위치를 이용해 구간의 최소 크기를 계속 업데이트한다. 이러한 방식으로 문제를 풀면 $O(n \log n)$ 의 시간복잡도로 문제를 통과할 수 있다.

하지만 더 효율적인 방법이 있다. 이를 위해서는 이 문제에서 모든 수열의 원소가 '자연수'라는 조건에 주목해야 한다. 수열의 원소가 모두 자연수라면, 구간의 길이를 늘릴수록 합은 무조건 커지고(단조 증가), 줄일수록 합은 작아진다. 이 성질을 이용하면 굳이 지난 합을 기억할 필요 없이, 두 개의 포인터(l, u)를 조절하는 것만으로 $O(N)$ 만에 답을 구할 수 있다.

🔗 문제 링크

코드

필요한 변수를 선언한다. 구간합의 가능한 최대 길이보다 5 큰 값으로 초기화하였다. 이는 구간 길이의 최소값을 업데이트 할 때 필요하다.

static int MAX_LENGTH = 100005;

다음으로 main 메서드에서 N과 S값을 입력받는다. N은 수열의 길이이고, S는 앞서 말했듯 구간합 크기의 기준이다.

public static void main(String[] args) throws IOException {
    StringTokenizer st = new StringTokenizer(br.readLine());
    int N, S;
    N = Integer.parseInt(st.nextToken());
    S = Integer.parseInt(st.nextToken());

수열의 원소를 입력받으면서 즉시 누적합 계산을 하여 prefixSum 배열에 저장한다. 인덱싱은 1 based로 하였다.

    int[] prefixSum = new int[N+1];
    st = new StringTokenizer(br.readLine());
    for (int i = 1; i <= N; i++) {
        prefixSum[i] = prefixSum[i-1] + Integer.parseInt(st.nextToken());
    }

이제 투 포인터 알고리즘을 수행한다. l은 왼쪽 포인터, u는 오른쪽 포인터를 나타낸다. minLength는 S보다 크거나 같은 조건을 만족하는 구간합의 길이의 최소값을 추적하기 위한 변수이다. 아까 언급했듯이, 초기화를 배열의 길이보다 5 큰 값으로 하여 최소값으로 업데이트가 제대로 이루어질 수 있도록 하였다.

가능한 모든 구간합을 확인하기 위해 반복문을 u <= N일 때까지 반복한다. 누적합을 통해 (l,u]의 구간합을 구한다. 그런 다음 그 구간합이 S보다 크거나 같다면 값을 줄이기 위해 왼쪽 포인터를 앞으로 이동시킨다. 만약 반대로 구간합이 S보다 작다면, 값을 늘리기 위해 오른쪽 포인터를 앞으로 이동시킨다.

    int l = 0;
    int u = 1;
    int minLength = MAX_LENGTH;
    while (u <= N) {
        int sectionSum = prefixSum[u]-prefixSum[l];
        if (sectionSum >= S) {
            minLength = Math.min(minLength, u-l);
            l++;
        }
        else u++;
    }

반복문을 빠져나왔다면 결과를 출력한다. 만약 minLength == MAX_LENGTH라면 한 번도 구간합의 값이 S보다 크거나 같지 않았던 것이므로 문제의 요구에 따라 0을 출력한다. 그렇지 않다면 minLength를 출력해 가능한 최소 길이를 출력한다.

    System.out.println(minLength == MAX_LENGTH ? 0 : minLength);
}

전체 코드

import java.util.*;
import java.io.*;

class Main {
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

    static int MAX_LENGTH = 100005;

    public static void main(String[] args) throws IOException {
        StringTokenizer st = new StringTokenizer(br.readLine());
        int N, S;
        N = Integer.parseInt(st.nextToken());
        S = Integer.parseInt(st.nextToken());

        int[] prefixSum = new int[N+1];
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= N; i++) {
            prefixSum[i] = prefixSum[i-1] + Integer.parseInt(st.nextToken());
        }

        int l = 0;
        int u = 1;
        int minLength = MAX_LENGTH;
        while (u <= N) {
            int sectionSum = prefixSum[u]-prefixSum[l];
            if (sectionSum >= S) {
                minLength = Math.min(minLength, u-l);
                l++;
            }
            else u++;
        }
        System.out.println(minLength == MAX_LENGTH ? 0 : minLength);
    }
}

시간복잡도

투 포인터에서, $u=N$이 되고 $l=N-1$이 되는 최악의 경우를 생각해볼 수 있다. 가령 다음 입력과 같은 경우이다.

10 10000
1 2 3 4 5 6 7 8 9 10000

즉 최악의 경우 투 포인터는 $O(2N-1) \approx O(N)$ 시간복잡도가 된다. 따라서 이 알고리즘은 BST를 사용한 $O(N \log N)$ 보다 빠르다.