에라토스테네스의 체를 비트마스킹으로 최적화 해보자.(Python, C++)

이런 의문이 들었어요.

제가 이번에 업로드한 에라토스테네스의 체를 보면 소수 리스트 구하는 과정에서 오직 True, False만으로 구현되어 있는 것을 보게 됩니다.

 

그리고 True와 False는 각각 1과 0이겠죠?

그래서 저는 이런 생각이 들었어요.

 

"이거 잘 생각해보니 bitset 형식으로 저장하면 되지 않을까?"

 

예시로 n = 7이라고 가정합니다.

그러면 bit = (1 << (7 + 1)) - 1 연산을 해줍니다.

bit = 0b11111111이 됩니다.

여기서 0과 1은 소수가 아니므로 0b11111100이 됩니다.

이 때 0b11111111 & ~(0b00000011) = 0b11111111 & 0b11111100 = 0b11111100이니 2^0 자리와 2^1이 0이 된 것을 볼 수 있습니다.

그리고 에라토스테네스의 체 알고리즘에 따라 2^2 자리수와 2^3 자리수는 1이 되고 2^4, 2^6은 0으로 만들어주는 식으로 하면

bit = 0b10101100

이렇게 되는데 정확히 2^p(여기서 p는 소수) 자리수만 1이 됩니다.

 

하지만 low-level에서 보면 ~ 연산도 0과 1을 반전시켜야 하는 시간이 걸리게 되니 여기서 역발상을 합니다.

1을 소수가 아닌 것으로, 0을 소수인 것으로 판정합니다.

이러면 어떻게 되는지 아래에 확인합니다.

 

n = 7이라고 가정합니다.

그런데 bit = (1 << (7 + 1)) - 1 연산을 하지 않아도 그냥 bit = 3하면 됩니다.

그러면 0b00000011이 되고 자동으로 2^0, 2^1 자리수는 0이 되니 0과 1은 소수가 아닌 것으로 판정됩니다.

 

그 다음 2과 3은 소수가 맞으니 2^2, 2^3 자리수는 0이 되고

그 다음 4, 6은 bit |= 1 << 4, bit |= 1 << 6 연산만 하면 됩니다.

index = 1 << 4로 가정합니다. 그러면 bit와 index는 다음과 같이 됩니다.

0b00000011(bit)

0b00010000(index)

bit |= index 연산해주면 bit = 0b00010011이 됩니다.

그 다음 index <<= 2 연산 해주면 비트 2칸만 시프트 하면 되니 엄청 빠른 최적화가 가능합니다.

그렇게 해서

0b00010011(bit)

0b01000000(index)

이 둘을 or 연산하면 bit = 0b01010011이 됩니다.

역시 2^p(p는 소수)만 0이 되니 bitset + bit mask 알고리즘으로 low-level에서 에라토스테네스의 체를 구현할 수 있습니다.

 

이것을 Python으로 구현합니다.

import math
import time

def bitset_sieve(n:int) -> int :
    if n < 2 :
        return 3 # 2 미만이면 0과 1만 합성수(0b11 = 3) 처리
        
    # 초기식 : 0b11 (0번, 1번 비트를 1[합성수]로 마킹)
    bit = 3 
    
    # 상한선 비트 마스크 사전 정의
    limit = 1 << n

    # 1. 짝수 폭격 (4부터 시작해서 비트 자릿수를 2씩 늘려감)
    index = 1 << 4
    while index <= limit :
        bit |= index
        index <<= 2 # 비트 자릿수를 +2 하는 효과 (즉, 4, 6, 8, 10...)
        
    # 2. 홀수 체 가동 (i = 3부터 시작)
    index_i = 1 << 3
    limit_sqrt = 1 << (int(math.sqrt(n)) + 1)
    
    # i를 숫자로 바꾸지 않고 비트 매커니즘만으로 j 추적하기
    index_j_start = 1 << 9   # i=3 일 때의 시작점 i^2 = 9
    jump_to_next_square = 16 # 다음 소수의 제곱자리로 워프할 크기 (5^2 - 3^2 = 16)
    inner_step = 6           # 내부 루프에서 점프할 비트 자릿수 크기 (2 * 3 = 6)
    
    while index_i < limit_sqrt :
        # 역발상 : bit와 AND 연산했을 때 0이 나와야 소수(True)입니다.
        if not (bit & index_i) :
            
            # i*i 비트 마스크 위치부터 n까지 폭격 시작
            index_j = index_j_start
            while index_j <= limit :
                bit |= index_j
                index_j <<= inner_step # 비트 자릿수를 +inner_step 만큼 밀어줌
                
        # 다음 홀수(i += 2) 단계의 파라미터들을 쉬프트 연산으로 갱신
        index_i <<= 2                         # i가 2 증가하므로 index_i는 2비트 전진
        index_j_start <<= jump_to_next_square # (i+2)^2 자리로 정확하게 마스크 워프
        
        # 수학적 불변성에 의해 다음 스퀘어 점프 거리는 8씩, 내부 스텝은 4씩 증가합니다.
        jump_to_next_square += 8  # (i+2)^2 - i^2 의 차이는 매동 8씩 증가함
        inner_step += 4           # 2*(i+2) - 2i = 4 이므로 간격은 항상 4씩 증가함
        
    return bit

# 검증용 프린트 함수 (0으로 살아남은 비트 인덱스만 추출)
def print_primes_from_bitset(bit_result:int, n:int) :
    primes = []
            
    index_number:int = 2
    index_shift:int = 1 << 2
    
    while index_number <= n :
        if not (bit_result & index_shift) :
            primes.append(index_number)
        index_number += 1
        index_shift <<= 1
    
    return primes

N:int = 1000000
TIME = 0.0
REPEAT = 1
for _ in range(REPEAT) :
    start_time = time.time()
    result_bit = bitset_sieve(N)
    result:list = print_primes_from_bitset(result_bit, N)
    end_time = time.time()
    TIME += end_time - start_time
    
print(f"경과 시간 : {TIME / REPEAT:.6f} 초")

테스트를 해봤으나 결과가 나오긴 하나 너무 오래 걸려서 실패합니다.

 

그 이유를 보니 Python이 자동으로 객체 할당을 지속적으로 해주는데 거기서 문제가 발생한 것입니다.

즉, 아주 큰 정수를 메모리 할당/해제를 계속 반복해줘서 문제가 발생합니다.

 

그러므로 bytearray(원시 데이터 배열) 혹은 array('Q')(64비트 정수 배열)을 사용해서 이를 해결하도록 하겠습니다.

import math
import time
from array import array

def bitset_sieve_array(n: int) -> array:
    if n < 2 :
        # 2 미만이면 0번, 1번만 합성수(0b11)로 켠 배열 반환
        L = array('Q', [0])
        L[0] = 3
        return L
        
    # [핵심] 64비트(8바이트) 정수 단위로 배열 '정적 할당'
    # n개의 비트를 담으려면 (n // 64) + 1 개의 64비트 메모리 블록이 필요함
    size = (n >> 6) + 1 
    bit_array = array('Q', [0] * size)
    
    # 0번, 1번 비트는 합성수(1)로 초기화
    bit_array[0] |= 3
    
    # 1. 짝수를 합성수(1)로 초기화 (4, 6, 8 ...)
    for j in range(4, n + 1, 2):
        # j >> 6 은 배열의 블록 위치, j & 63 은 64비트 내부의 정확한 자릿수 타격
        bit_array[j >> 6] |= (1 << (j & 63))
        
    # 2. 홀수 체 가동 (i = 3부터)
    limit_sqrt = int(math.sqrt(n)) + 1
    
    for i in range(3, limit_sqrt, 2) :
        # i번째 비트가 0(소수)인지 확인
        if not (bit_array[i >> 6] & (1 << (i & 63))) :
            
            # i*i부터 2*i 간격으로 점프하며 합성수 타격 (이전 동일한 2점프 알고리즘)
            for j in range(i * i, n + 1, i << 1) :
                # 파이썬 메모리 복사 없이, C언어 레벨에서 해당 비트만 O(1) 수정!
                bit_array[j >> 6] |= (1 << (j & 63))
                
    return bit_array

# 검증용 프린트 함수 (사용자님의 연속 시프트 아이디어 + 64비트 청크 최적화)
def print_primes_from_array(bit_array: array, n: int) -> list:
    primes = []
    
    word_index = 0
    bit_shift = 1 << 2 # 2번 비트부터 검사 시작
    
    index = 2
    # 'index_shift <<= 1' 방식 유지
    while index <= n :
        if not (bit_array[word_index] & bit_shift):
            primes.append(index)
            
        bit_shift <<= 1 # 단 1비트만 밀어내며 광속 전진
        
        # 64비트 경계를 넘어가면, 다음 배열 블록으로 넘어가고 시프트 초기화
        if bit_shift == (1 << 64):
            bit_shift = 1
            word_index += 1
        
        index += 1
            
    return primes

N = 1000000
TIME = 0.0
REPEAT = 10

for _ in range(REPEAT):
    start_time = time.time()
    result_L = bitset_sieve_array(N)
    result_list = print_primes_from_array(result_L, N)
    end_time = time.time()
    TIME += end_time - start_time
    
print(f"경과 시간 : {TIME / REPEAT:.6f} 초")

실행해본 결과 다음과 같습니다.

약 4.5초가 나옵니다.

원래 에라토스테네스의 체 소스 코드가 0.1초 나온 것을 보면 오히려 시간이 더 늘어난 것을 확인할 수 있습니다.

이것은 Python 내부에서 integer 객체는 불변 객체라서 강제로 자동할당 해줍니다.

이게 병목 구간인 것으로 보입니다.

 

그 외 제가 Python으로 구현을 아무리 해도 Python의 태생적인 한계(강제로 정수 객체를 재할당)에 의해 Python은 포기하고 그 대신 C++로 변경합니다.

 

#include <iostream>
#include <vector>
#include <cmath>

using namespace std;

// 소수 리스트를 반환하는 함수
vector<int> eratosthenes_sieve(int n) {
    if (n < 2) return {};

    // vector<bool>은 비트 단위로 압축되어 메모리를 극도로 절약합니다.
    vector<bool> is_prime(n + 1, true);
    
    is_prime[0] = is_prime[1] = false;

    // 1. 짝수 처리 (4부터 n까지, 2씩 증가)
    for (int i = 4; i <= n; i += 2) {
        is_prime[i] = false;
    }

    // 2. 홀수 처리 (3부터 sqrt(n)까지, 2씩 증가)
    int limit = static_cast<int>(std::sqrt(n));
    for (int i = 3; i <= limit; i += 2) {
        if (is_prime[i]) {
            // i*i부터 2*i 간격으로 배수 제거
            for (int j = i * i; j <= n; j += 2 * i) {
                is_prime[j] = false;
            }
        }
    }

    // 3. 소수 수집
    vector<int> primes;
    primes.reserve(n / 10); // 성능 최적화 : 예상 크기만큼 미리 메모리 할당
    for (int i = 2; i <= n; i++) {
        if (is_prime[i]) {
            primes.push_back(i);
        }
    }

    return primes;
}

int main() {
    int n = 100;
    vector<int> result = eratosthenes_sieve(n);

    // 결과 출력
    cout << "[";
    for (size_t i = 0; i < result.size(); i++) {
        cout << result[i] << (i == result.size() - 1 ? "" : ", ");
    }
    cout << "]" << '\n';

    return 0;
}

이것이 바로 기존에 있던 Python 소스 코드를 C++로 변경한 것입니다.

여기에 시간 측정하는 코드도 추가할게요.

#include <iostream>
#include <vector>
#include <cmath>
#include <chrono> // 시간 측정을 위한 헤더 추가

using namespace std;

// 소수 리스트를 반환하는 함수
vector<int> eratosthenes_sieve(int n) {
    if (n < 2) return {};

    // vector<bool>은 비트 단위로 압축되어 메모리를 극도로 절약합니다.
    vector<bool> is_prime(n + 1, true);
    
    is_prime[0] = is_prime[1] = false;

    // 1. 짝수 처리 (4부터 n까지, 2씩 증가)
    for (int i = 4; i <= n; i += 2) {
        is_prime[i] = false;
    }

    // 2. 홀수 처리 (3부터 sqrt(n)까지, 2씩 증가)
    int limit = static_cast<int>(std::sqrt(n));
    for (int i = 3; i <= limit; i += 2) {
        if (is_prime[i]) {
            // i*i부터 2*i 간격으로 배수 제거
            for (int j = i * i; j <= n; j += 2 * i) {
                is_prime[j] = false;
            }
        }
    }

    // 3. 소수 수집
    vector<int> primes;
    primes.reserve(n / 10); // 성능 최적화 : 예상 크기만큼 미리 메모리 할당
    for (int i = 2; i <= n; i++) {
        if (is_prime[i]) {
            primes.push_back(i);
        }
    }

    return primes;
}

int main() {
    int n = 1000000; // 범위 지정
    int repeat = 100; // 반복 횟수
    double total_duration = 0.0;

    for (int i = 0; i < repeat; i++) {
        // 측정 시작
        auto start = chrono::high_resolution_clock::now();
        
        vector<int> result = eratosthenes_sieve(n);
        
        // 측정 종료
        auto end = chrono::high_resolution_clock::now();
        chrono::duration<double> elapsed = end - start;
        
        total_duration += elapsed.count();
    }

    // 평균 경과 시간을 출력합니다.
    cout << "average elapsed time : " << total_duration / repeat << " second" << endl;

    return 0;
}

C++에서 기존 에라토스테네스의 체 알고리즘 기반으로 실행해본 결과 다음과 같이 나왔습니다.

약 0.045 초가 나옵니다.

이제 bit mask 방식을 적용해서 최적화를 시도합니다.

#include <iostream>
#include <vector>
#include <cmath>
#include <chrono>

using namespace std;

using ull = unsigned long long;

// 0 = 소수(True), 1 = 합성수(False)
vector<ull> bitmask_sieve_cpp(ull n) {
    if (n < 2) return {};

    // 64비트(8바이트) 단위로 배열 할당. 초기값 0ULL(모두 소수)
    ull size = (n >> 6) + 1;
    vector<uint64_t> bit_array(size, 0ULL);

    // 0번, 1번 비트는 합성수(1)로 켜기 (0b11 = 3)
    bit_array[0] |= 3ULL;

    // 1. 짝수 폭격 (4, 6, 8 ...)
    for (ull j = 4; j <= n; j += 2) {
        // j >> 6 은 블록 인덱스, j & 63 은 0~63 사이의 비트 위치
        bit_array[j >> 6] |= (1ULL << (j & 63));
    }

    // 2. 홀수 체 가동 (i = 3부터)
    ull limit = static_cast<ull>(sqrt(n));
    for (ull i = 3; i <= limit; i += 2) {
        // i번째 비트가 0(소수)인지 확인 (AND 연산 결과가 0이면 소수)
        if (!(bit_array[i >> 6] & (1ULL << (i & 63)))) {
            
            // i*i부터 2*i 간격으로 합성수 타격
            ull step = i << 1;
            for (ull j = i * i; j <= n; j += step) {
                bit_array[j >> 6] |= (1ULL << (j & 63));
            }
        }
    }

    // 3. 소수 수집 (연속 시프트 최적화 기법 적용)
    vector<ull> primes;
    primes.reserve(n / 10);
    
    uint64_t bit_shift = 1ULL << 2; // 2번 비트부터 스캔 시작
    ull word_index = 0;
    
    for (ull i = 2; i <= n; i++) {
        // 비트가 0이면 소수로 판정하여 추가
        if (!(bit_array[word_index] & bit_shift)) {
            primes.push_back(i);
        }
        
        bit_shift <<= 1; // 1비트 전진 (O(1))
        
        // C++에서 uint64_t의 1을 64번 밀어내면 왼쪽으로 빠져나가 0이 됨
        // 0이 되는 순간 다음 64비트 블록으로 넘어가고 마스크 초기화
        if (bit_shift == 0) {
            bit_shift = 1ULL;
            word_index++;
        }
    }

    return primes;
}

int main() {
    ull n = 1000000000;  // 범위 지정
    int repeat = 10; // 반복 횟수
    double total_duration = 0.0;

    for (int i = 0; i < repeat; i++) {
        // 측정 시작
        auto start = chrono::high_resolution_clock::now();
        
        vector<ull> result = bitmask_sieve_cpp(n);
        
        // 측정 종료
        auto end = chrono::high_resolution_clock::now();
        chrono::duration<double> elapsed = end - start;
        
        total_duration += elapsed.count();
    }

    // 평균 경과 시간을 출력합니다.
    cout << "average elapsed time : " << total_duration / repeat << " second" << endl;

    return 0;
}

여기서 n = 10^9, repeat = 10으로 테스트를 진행합니다.

n = 10^9 이하의 소수들을 겨우 약 19.3초 만에 모두 찾았군요!

그 아래에 n = 10^6, repeat = 1000으로 테스트를 진행합니다.

    ull n = 1000000;  // 범위 지정
    int repeat = 1000; // 반복 횟수

이러면 결과가 어떻게 나올까요?

정말로 말도 안 되는 결과가 나옵니다.

n = 10^6이면 그래도 나름대로 큰 수인데요.

이게 겨우 12[ms]만에 바로 10^6 이하의 소수들을 모두 구해냅니다.

자, 이제 n = 10^3, repeat = 100000로 하면 어떻게 나올까요?

    ull n = 1000;  // 범위 지정
    int repeat = 100000; // 반복 횟수

이러면 결과가 어떻게 나올까요?

약 14[μs]가 나옵니다.

 

이번 연구 활동을 통해 Python은 태생적인 한계로 인해 막혔으나 C++에서 에라토스테네스의 체 알고리즘에서 좀 더 최적화 하는데 성공했습니다!

 

참고로 이 소스 코드는 unsigned long long 범위로 지정해서 n이 최대 2^64 - 1 이하의 소수들을 모두 찾아줍니다.

이론상으로 10^18까지 구할 수 있으나 그렇게 하면 최적화된 소스 코드로도 모든 소수들을 구하는데 약 수백년 정도 걸려서 현실적으로 볼 때 10^12으로 하면 기존 에라토스테네스의 소스 코드로는 하루 종일 걸리는 것이 최적화 된 소스 코드를 활용하면 5~6시간 만에 10^12 이하의 모든 소수들을 다 구해줘서 10^24 이하의 자연수의 소수 판정이 가능해집니다.

 

오늘 연구 성과는 대성공입니다.

기존 에라토스테네스의 체 알고리즘에서 bitset + bit mask 알고리즘으로 동일한 메모리 공간에서 경과 시간을 약 4배 정도 최적화 하는데 성공했습니다.