[코딩테스트] [Two Pointer] [2003] 수들의 합 2 & [1806] 부분합

투 포인터 알고리즘

코딩테스트에서 자주 나오는 유형 중 하나인 투 포인터 알고리즘에 대해 알아보자.

먼저 말로 풀어보면 선형 데이터 구조에서 두 개의 인덱스를 관리하여 특정 조건을 만족하는 부분집합이나 특정값이다.

이렇게 해서는 이해하기 어려우니 문제 예시를 보며 이해를 해보자.

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

출력

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

문제풀이

이 문제는 N개의 자연수 수열의 i~j번째 수의 합이 M이 되는 경우를 구하는 문제이다.

[i:j]합이 M이 된다면 count를 1증가 시키면 된다.

Brute Force

가장 쉽게 생각해서 할 수 있는 방법은 모든 경우를 검사하는 것이다.

아래 코드를 보면 이중 반복문을 사용해서 모든 요소를 검사하여 Sum이 M이 될 때 break한다.

int ans = 0;
for (int i = 0; i < N; i++) {
	int sum = 0;
	for (int j = i; j < N; j++) {
		sum += arr[j];
		if (sum == M) ans++;
		if (sum >= M) break;
	}
}

하지만 이 방법은 시간 복잡도가 O(N^2)이므로 시간 효율이 좋지 않다.

Binary Search + 누적합

자연수 수열이어서 가능한 방법이 있다.

각 시작점 i에 대해 [i:j]가 M이 될 수 있는지 Binary Search + 누적합으로 해결할 수 있다.

int ans = 0;
for (int i = 1; i <= N; i++) {
	int l = i, r = N;
	while (l <= r) {
		int m = (l + r)/2;
		long sum = acc[m] - acc[i-1];
		if (sum > M) r = m -1;
		els if (sum < M) l = m + 1;
		else {
			ans++;
			break;
		}
	}
}

Two Pointer

[i:j]의 j가 늘어나면 구간의 합이 증가한다는 특성을 이용해 [i:j] == M을 만족하는 구간을 Two Pointer로 찾아보자.

이 원리는 i번째에서 j번째까지 합을 구할 때 합이 M보다 크다면 i를 옮겨서 값을 줄이고 M보다 작다면 j를 옮겨서 늘리는 방식이다.

만약 값이 같다면 count를 증가해주면 된다.

즉, 각 i에 대해 j를 이동해 [i:j]의 합이 M이 되는 구간을 찾는다.

1. 현재 구간의 합이 M보다 크거나 같을때까지 j를 이동시킨다.
2. 합이 M이라면 답을 증가시킨다.
3. 이후 i를 이동하므로 현재 구간의 합에서 A[i]를 제한다.

int rightIndex = 0;
int currentSum = arr[0];
for (int i = 0; i < N; i++) {
	while (currentSum < M && rightIndex < N - 1)
		currentSum += arr[++rightIndex];
	if (currentSum == M) ans++;
	currentSum -= arr[i];
}

전체코드

import java.util.Scanner;

class Main
{
    public static void main (String[] args) {
        Scanner sc = new Scanner(System.in);

        int N = sc.nextInt();
        int M = sc.nextInt();
        int[] arr = new int[N];
        for (int i = 0; i < N; i++)
            arr[i] = sc.nextInt();

        int ans = 0;
        int nextIndex = 0;
        int currentSum = 0;
        for (int i = 0; i < N; i++) {
            while (currentSum < M && nextIndex < N)
                currentSum += arr[nextIndex++];
            if (currentSum == M) ans++;
            currentSum -= arr[i];
        }

        System.out.println(ans);
    }
}

시간복잡도

이 알고리즘은 i가 0부터 N-1까지 움직이는동안 rightIndex의 총 움직임이 N이므로 전체 시간복잡도는 O(N2)이 아닌 O(N)이 된다.

[1806] 부분합

시간 제한 메모리 제한 제출 정답 맞힌 사람 정답 비율

0.5 초 (하단 참고)

128 MB

113642

31730

22449

26.346%

문제

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

입력

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

출력

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

문제풀이

이 문제는 N개의 자연수 수열에서 연속된 수들의 부분합 중에 그 합이 S이상이 되는 것 중 가장 짧은 것의 길이를 구하는 문제이다.

N이 10만이라서 O(N^2)이 나오면 시간초과가 발생한다.

예제 1를 보면 부분합중 15 이상이 되는것중 최소의 길이를 구하는 것이다.

아래 이미지는 15이상이 되는 부분합을 나열한 이미지다.

각 left index에 대해 가장 짧은 길이를 구하는 방법은 두가지로 할 수 있다.

BinarySearch + 누적합

자연수 수열이므로 각 시작점 i에 대해 [i:j]가 S이상이 될 수 있는지 해결할 수 있다.

import java.util.Scanner;

class Main
{
    public static void main (String[] args) {
        Scanner sc = new Scanner(System.in);

        int N = sc.nextInt();
        int S = sc.nextInt();
        int[] arr = new int[N + 1];
        for (int i = 1; i <= N; i++)
            arr[i] = sc.nextInt();

        long[] acc = new long[N + 1];
        for (int i = 1; i <= N; i++)
            acc[i] = acc[i - 1] + arr[i];

        int ansLength = N + 1;
        for (int i = 1; i <= N; i++) {
            int l = i, r = N;
            while (l <= r) {
                int m = (l + r) / 2;
                long sum = acc[m] - acc[i - 1];
                if (sum >= S) {
                    ansLength = Math.min(ansLength, m - i + 1);
                    r = m - 1;
                }
                else l = m + 1;
            }
        }

        if (ansLength > N) ansLength = 0;
        System.out.println(ansLength);
    }
}

시간복잡도

이 알고리즘의 시간 복잡도는 O(NlogN)이 될 것이다.

Two Pointer

l1 < l2의 가장 짧은 구간 [l1:r1], [l2:r2]에서 항상 r1 ≤ r2이다.

따라서 left index를 차례로 순회하며 대응하는 right index를 증가하는 방향으로 유지하며 투포인터를 사용할 수 있다

중요한 부분은 현재 구간의 합이 S보다 작을때는 계속 구간을 늘려 합을 증가시켜주자.

이 때, edge case를 놓치면 안된다. 만약 nextIndex가 N이면 더 이상 계산이 불가하다.

구간의 합이 S 이상이면 조건을 만족한 것이고 i를 증가 시켜 최소의 길이를 가진 구간합을 찾으면 된다.

        int ansLength = N + 1;
        int nextIndex = 0;
        int currentSum = 0;
        for (int i = 0; i < N; i++) {
            while (currentSum < M && nextIndex < N)
                currentSum += arr[nextIndex++];
            if (currentSum >= M) ansLength = Math.min(ansLength, nextIndex - i);
            currentSum -= arr[i];
        }

전체코드

import java.util.Scanner;

class Main
{
    public static void main (String[] args) {
        Scanner sc = new Scanner(System.in);

        int N = sc.nextInt();
        int M = sc.nextInt();
        int[] arr = new int[N];
        for (int i = 0; i < N; i++)
            arr[i] = sc.nextInt();

        int ansLength = N + 1;
        int nextIndex = 0;
        int currentSum = 0;
        for (int i = 0; i < N; i++) {
            while (currentSum < M && nextIndex < N)
                currentSum += arr[nextIndex++];
            if (currentSum >= M) ansLength = Math.min(ansLength, nextIndex - i);
            currentSum -= arr[i];
        }

        if (ansLength > N) ansLength = 0;
        System.out.println(ansLength);
    }
}

시간복잡도

이 알고리즘의 시간복잡도는 O(2N)이다.

즉, O(N)이다.