[백준 5615] 아파트 임대 (Feat. 밀러-라빈 소수 판정법) (Platinum I)

https://www.acmicpc.net/problem/5615

5615번: 아파트 임대

 

매우 자명한 어떤 절차에 의해, 2N + 1이 소수가 되는 N을 찾는 문제임을 알 수 있다. Naive한 소수 판별을 짜면 시간복잡도 \(O(T \sqrt{N})\)이 나와서 TLE가 난다. 시간복잡도를 줄일 수 있는 방법으로 밀러-라빈 소수 판정법을 이용하자.

a = 2, 7, 61에 대하여

1. \( n - 1 = 2^s d \)로 놓는다.
2. \( a^d \ne 1 \mod n \)이고,
3. 모든 r : 0 ~ s-1에 대하여 \( a^{2^r d} \ne -1 \mod n \)이면 n은 합성수이다.

이 절차는 n < 48억인 모든 n에 대하여 성립한다. 더 큰 n에 대해서는 a값을 다르게 잡아줘야 한다.

• a = 2, 3, 5, 7, 11로 잡으면 n < 2조인 모든 n에 대하여 성립한다.
• a = 2, 3, 5, 7, 11, 13, 17로 잡으면 n < 340조인 모든 n에 대하여 성립한다.

이렇게 더 큰 숫자일수록 n을 더 많이 잡아야 하는 이유는, 밀러-라빈 소수 판정법이 원래는 확률적 알고리즘이었기 때문이다. 그러니까 밀러-라빈 소수 판정법은 원래 합성수를 판별하는 알고리즘이다. 즉, a를 아무렇게나 잡으면, 운이 좋다면 n이 합성수라는 결론을 얻어낼 수는 있지만, n이 소수라는 결론을 얻어낼 수는 없다. 이 알고리즘은 한 번 실행할 때마다 합성수의 3/4가 걸러진다는 사실이 알려져 있으므로, 서로 다른 10개쯤의 a에 대해서 n이 이 알고리즘을 통과하면 그냥 n은 소수라고 판정해 버리는 것이다.

시간복잡도는 한 번 돌릴 때마다 \( O(k \log^3 n) \). k는 a의 개수이다.

#include <bits/stdc++.h>

enum RESULT {COMPOSITE = 0, MAYBE_PRIME};

auto modulo_power_factory(unsigned long long n) {
    std::function<unsigned long long(unsigned long long, unsigned long long)> pow = 
            [&pow, n](unsigned long long x, unsigned long long p) -> unsigned long long {
        if (p <= 1) {
            return p <= 0 ? 1 : x % n;
        }
        auto temp = pow(x, p / 2);
        if (p % 2 == 0) {
            return temp * temp % n;
        } else {
            return temp * temp % n * x % n;
        }
    };

    return pow;
}

auto miller_rabin_factory(int a) {
    return [a](unsigned long long n) -> RESULT {
        if (n % 3 == 0 and n > 3) {
            return COMPOSITE;
        }
        
        if (n <= a) {
            return MAYBE_PRIME;
        }
        int s = std::countr_zero(static_cast<unsigned long long>(n - 1));
        unsigned long long d = (n - 1) / (1LL << s);
        
        auto modulo_power_n = modulo_power_factory(n);

        if (modulo_power_n(a, d) == 1) {
            return MAYBE_PRIME;
        }
        for (int i = 0; i < s; ++i) {
            if (modulo_power_n(a, d * (1LL << i)) == n - 1) {
                return MAYBE_PRIME;
            }
        }

        return COMPOSITE;
    };
}

void solve() {
    int t; std::cin >> t;
    int ret = 0;
    while (t --> 0) {
        unsigned long long x; std::cin >> x;
        x = 2 * x + 1;

        auto res = MAYBE_PRIME;
        for (int a : {2, 7, 61}) {
            if (miller_rabin_factory(a)(x) == COMPOSITE) {
                res = COMPOSITE;
                break;
            }
        }
        ret += (res == MAYBE_PRIME ? 1 : 0);
    }

    std::cout << ret << std::endl;
}

#ifndef LOCAL
int main() {
    std::cin.tie(nullptr)->sync_with_stdio(false);
    solve();
}
#endif