1644 - 소수의 연속합

이 문제는 입력으로 어떤 자연수가 주어졌을 때, 연속된 소수의 합을 이용해 그 자연수를 나타낼 수 있는 경우의 수가 몇 개인지 출력하는 코드를 요구하고 있다.

여기에서 주목할 점은 주어지는 N의 범위가 최대 4백만으로 매우 크다는 것이다. 즉 시간복잡도가 $O(N)$에서 크게 벗어나지 않는 코드를 작성해야 한다. 그렇다면 어떻게 해야 이 제약조건 하에서 문제를 풀 수 있을까?

여기서 우리는 소수의 "연속합" 부분에 주목해야 한다. N 이하의 소수들에 대한 리스트가 있다고 할 때, 이 리스트에서 두 지점(시작과 끝)을 잡아 그 사이에 있는 모든 숫자를 더해본다고 생각하자. 그렇다면 끝 지점을 오른쪽으로 옮기면 그 구간의 합의 크기는 반드시 증가할 것이고, 시작 지점을 오른쪽으로 옮기면 그 구간의 합의 크기는 반드시 감소할 것이다. 따라서 투 포인터 알고리즘을 이용해 소수 수열을 탐색하는 방식으로 이 문제를 풀어볼 수 있다.

🔗 문제 링크

코드

먼저 연속된 소수의 합으로 나타낼 수 있는 경우의 수를 구할 숫자 N을 입력받는다. N이 1이라면, 1은 소수도 합성수도 아니므로 0을 출력해준다.

public static void main(String[] args) throws IOException {
    int N = Integer.parseInt(br.readLine());

    if (N == 1) {
        System.out.println(0);
        return;
    }

다음으로 N 이하의 소수 수열을 구하기 위해 에라토스테네스의 체 알고리즘을 이용해 소수를 판정한다. 소수 여부를 저장해두기 위해 isPrime 배열을 선언하고, N이하의 소수들을 저장해두기 위해 primes 배열을 선언한다. pnum은 N 이하 소수의 개수를 추적하기 위한 변수이다.

isPrime 배열의 모든 원소를 true로 초기화한 뒤, 에라토스테네스 체 알고리즘을 이용해 소수를 판정하고, 소수의 배수에 해당하는 수에 대해 대응되는 isPrime 원소의 값을 false로 만들어 소수가 아님을 표시해준다. 소수로 판정된 수는 primes 배열에 저장해준다.

    boolean[] isPrime = new boolean[N+1];
    int[] primes = new int[N+1];
    int pnum = 0;
    Arrays.fill(isPrime, true);

    for (int i = 2; i <= N; i++) {
        if (!isPrime[i]) continue;
        primes[pnum++] = i;
        for (long j = (long) i*i; j <= N; j += i) {
            isPrime[(int)j] = false;
        }
    }

이제 본격적으로 투 포인터 알고리즘을 수행해준다. 먼저 연속된 소수의 합으로 나타낼 수 있는 경우의 수를 추적하기 위해 cases 변수를 선언해준다. 또한 소수 수열에서 합을 구할 범위를 가리키는 s와 e 변수를 선언하고, 초기 위치는 모두 0으로 설정해 준다. 합하는 범위는 [s,e]로 생각하면 된다.

포인터가 가리키는 구간의 합을 추적하기 위해 sum 변수를 선언해 둔다. 초기에 포인터가 가리키는 범위는 [0, 0] 이므로 sum = primes[0]으로 초기화해 둔다.

다음으로 e < pnum일 때까지 반복문을 돌면서 소수 수열에서 합이 N인 연속 구간이 있는지 탐색한다.

sum == N이면 경우의 수가 발견된 것이므로 cases++를 통해 경우의 수를 늘려 준다.

sum >= N 이라면 구간합을 감소시킬 필요가 있으므로, 기존 시작 위치의 소수 값을 구간합 sum에서 빼준 뒤 s를 1 증가시킨다.

반대로 sum < N 이라면 구간합을 증가시킬 필요가 있으므로 e를 1 증가시킨 뒤 그 위치가 가리키는 소수 값을 구간합 sum에 추가해 준다. 참고로 e에 1을 더했을 때 pnum이 되면 탐색이 완료된 것이므로 break를 통해 반복문을 빠져나온다.

    int cases = 0;
    int s, e;
    s = e = 0;

    int sum = primes[0];


    while (e < pnum) {
        if (sum == N) cases++;

        if (sum >= N) {
            sum -= primes[s++];
        }
        else {
            if (e + 1 == pnum) break;
            sum += primes[++e];
        }
    }

반복문을 빠져나왔다면 경우의 수를 모두 구한 것이므로 cases 값을 출력하고 프로그램을 종료한다.

    System.out.println(cases);
}

전체 코드

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

class Main {
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    public static void main(String[] args) throws IOException {
        int N = Integer.parseInt(br.readLine());

        if (N == 1) {
            System.out.println(0);
            return;
        }

        boolean[] isPrime = new boolean[N+1];
        int[] primes = new int[N+1];
        int pnum = 0;
        Arrays.fill(isPrime, true);

        for (int i = 2; i <= N; i++) {
            if (!isPrime[i]) continue;
            primes[pnum++] = i;
            for (long j = (long) i*i; j <= N; j += i) {
                isPrime[(int)j] = false;
            }
        }
        int cases = 0;
        int s, e;
        s = e = 0;

        int sum = primes[0];


        while (e < pnum) {
            if (sum == N) cases++;

            if (sum >= N) {
                sum -= primes[s++];
            }
            else {
                if (e + 1 == pnum) break;
                sum += primes[++e];
            }
        }

        System.out.println(cases);
    }
}

시간복잡도

에라토스테네스의 체 시간복잡도는 $O(N \log \log N)$으로 알려져 있으며, 소수의 개수는 대략 $\frac{N}{\log N}$에 비례하므로, 투 포인터 계산 비용은 $O(\frac{N}{\log N})$ 이다.

두 항의 크기를 비교해 보면

$$\frac{N/\log N}{N \log \log N} = \frac{1}{\log N \log \log N} \to 0 \ (N \to \infty)$$

이므로 에라토스테네스의 체 비용이 더 큼을 알 수 있다. 따라서 위 코드의 시간복잡도는 $O(N \log \log N)$이다.