C++로 N 이하의 수 중 소수 개수와의 차이가 K 이상인 수의 개수 구하기

문제 정의

두 정수 N과 K가 주어졌을 때, 아래 조건을 모두 만족하는 수의 개수를 구하는 것이 목표입니다.

구하려는 수는 N 이하이어야 합니다.
그 수 이하에 존재하는 소수의 개수를 count라고 할 때, |수 − count| ≥ K를 만족해야 합니다.

예제 1

입력

N = 5, K = 2

출력

N 이하의 수 중 소수 개수와의 차이가 K 이상인 수의 개수: 2

설명
조건을 만족하는 수는 다음과 같습니다.
5 (5 − 2 ≥ 2)와 4 (4 − 2 ≥ 2)

예제 2

입력

N = 10, K = 6

출력

N 이하의 수 중 소수 개수와의 차이가 K 이상인 수의 개수: 1

설명
조건을 만족하는 수는 다음과 같습니다.
10 (10 − 4 ≥ 6)

접근 방법

이 문제는 이진 탐색(binary search)을 활용하면 연산량을 크게 줄일 수 있습니다. 핵심 아이디어는 다음과 같습니다.

num 이하의 소수 개수를 count1, num+1 이하의 소수 개수를 count2라고 하면, num+1 − count2 ≥ num − count1이 항상 성립합니다. 즉, 어떤 수가 조건을 만족하면 그보다 큰 수 역시 반드시 조건을 만족합니다. 따라서 이진 탐색으로 조건을 만족하는 가장 작은 수 num을 찾아내면, num부터 N 사이의 모든 수가 조건을 충족하므로 전체 개수를 한 번에 계산할 수 있습니다.

알고리즘 단계

N과 K를 입력받습니다.
배열 arr[]의 인덱스 i에는 i 이하의 소수 개수가 저장됩니다.
set_prime() 함수는 배열 arr[]에 소수 개수를 채워 넣는 역할을 합니다.
배열 check[i]는 i가 소수이면 true, 아니면 false를 저장합니다.
0과 1은 소수가 아니므로 check[0] = check[1] = false로 설정합니다.
i = 2부터 i * i < size(1000001)까지 check 배열을 순회하면서, check[i]가 1이면(소수이면) j = i * 2부터 size 미만까지 i씩 증가시키며 check[j]를 0으로 설정합니다. (에라토스테네스의 체)
이후 for 루프로 arr[]를 갱신합니다. 기본적으로 arr[i] = arr[i − 1]이며, i 자체가 소수라면 arr[i]++로 개수를 하나 늘립니다.
total(int N, int K) 함수는 N과 K를 받아, N 이하의 수 중 소수 개수와의 차이가 K 이상인 수의 개수를 반환합니다.
set_prime()을 호출한 뒤 temp_1 = 1, temp_2 = N, count = 0으로 초기화합니다.
이진 탐색 while 루프 안에서 set = (temp_1 + temp_2) >> 1, 즉 (첫 값 + 끝 값) / 2로 중간 지점을 정합니다.
set − arr[set] ≥ K이면 조건을 만족하므로 count를 set으로 갱신하고, temp_2 = set − 1로 탐색 범위를 앞쪽으로 줄입니다.
조건을 만족하지 않으면 temp_1 = set + 1로 탐색 범위를 뒤쪽으로 옮깁니다.
탐색이 끝나면 count는 조건을 만족하는 최소의 수가 되며, 최종 답은 N − count + 1입니다. 조건을 만족하는 수가 없다면 0이 됩니다.
모든 루프가 종료되면 count를 결과로 반환합니다.

C++ 구현 예제

#include <bits/stdc++.h>
using namespace std;
#define size 1000001
int arr[size];
void set_prime(){
    bool check[size];
    memset(check, 1, sizeof(check));
    check[0] = 0;
    check[1] = 0;
    for (int i = 2; i * i < size; i++){
        if(check[i] == 1){
            for (int j = i * 2; j < size; j += i){
                check[j] = 0;
            }
        }
    }
    for (int i = 1; i < size; i++){
        arr[i] = arr[i - 1];
        if(check[i] == 1){
            arr[i]++;
        }
    }
}
int total(int N, int K){
    set_prime();
    int temp_1 = 1;
    int temp_2 = N;
    int count = 0;
    while (temp_1 <= temp_2){
        int set = (temp_1 + temp_2) >> 1;
        if (set - arr[set] >= K){
            count = set;
            temp_2 = set - 1;
        } else {
            temp_1 = set + 1;
        }
    }
    count = (count ? N - count + 1 : 0);
    return count;
}
int main(){
    int N = 12, K = 5;
    cout << "N 이하의 수 중 소수 개수와의 차이가 K 이상인 수의 개수: " << total(N, K);
    return 0;
}

실행 결과

위 코드를 실행하면 다음과 같은 결과가 출력됩니다.

N 이하의 수 중 소수 개수와의 차이가 K 이상인 수의 개수: 4