[C++] 17104번 골드바흐 파티션 2

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

---

 

26/01/31

 

 

르모앙의 추측 문제를 약간 변형시킨 문제이다.

---

 

문제 접근 방식:

 

 

짝수 $N$을 두 소수 $p+q$로 나타내는 경우의 수를 세는 문제다.

얼핏보면 르모앙의 추측 문제를 해결했을 때와 똑같은 방법으로 해결할 수 있을 것이라고 생각할 수 있겠지만, 문제는 $p+q$와 $q+p$를 같은 경우의 수로 취급한다는 점이다.

Sieve를 곱한 결과는 $(p, q)$와 $(q, p)$를 다른 경우의 수로 취급한다.

따라서, 기본적으로 나온 결과에서 중복된 경우의 수를 빼줘야만 한다.

$N \neq 2p$인 경우, $p+q = N$을 만족시키는 $(p, q)$ 순서 쌍은 항상 $p\neq q$이다.

즉, Sieve끼리 곱한 결과(중복을 포함한 경우의 수)는 $(p, q)$와 $(q, p)$가 항상 쌍을 이루고 있다. 따라서, 절반으로 나눈것이 중복된 경우의 수를 뺀 것이다.

$N = 2p$인 경우, $p+q = N$을 만족시키는 $(p, q)$ 순서 쌍은 $(p, p)$를 제외하고 항상 $p\neq q$이다.

따라서 이 경우, Sieve끼리 곱한 결과(중복을 포함한 경우의 수)는 $(p, p)$를 제외하고 $(p, q)$와 $(q, p)$가 항상 쌍을 이루고 있으므로, $1$을 더하고 절반으로 나눈 것이 중복된 경우의 수를 뺀 것이다.

즉, 결과로 나온 계수 배열 $C$가 있다면 $(C[N]+1) / 2$가 답이 된다.

---

아래는 내가 위의 접근 방식과 같이 작성한 C++ 코드이다. 더보기를 누르면 확인할 수 있다.

// 17104번 골드바흐 파티션 2
// NTT
#include <iostream>
#include <vector>

using namespace std;
#define fastio ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
#define endl '\n'

template<long long MOD, long long G>
struct NTT {
    using ll = long long;
private :
    // HELPER FUNCTION
    inline int mul_mod(int a, int b) const {return (ll)a * b % MOD;}
    inline int pow_mod(int a, int e) const {
        int r = 1;
        while(e > 0){
            if(e & 1) r = mul_mod(r, a);
            a = mul_mod(a, a);
            e >>= 1;
        }
        return r;
    }
    inline int mod_inv(int a) const {return pow_mod(a, MOD-2);}
public : 
    void iterative_NTT(vector<int>& a, int invert){
        int n = (int)a.size();
        // 1) bit-reversal permutation
        for (int i = 1, j = 0; i < n; i++){
            int bit = n >> 1;
            while (j & bit){
                j ^= bit;
                bit >>= 1;
            }
            j ^= bit;
            if (i < j) swap(a[i], a[j]);
        }
        // 2) butterflies by length = 2, 4, 8, ...
        for (int len = 2; len <= n; len <<= 1){
            // primitive len-th root of unity
            int w_len = pow_mod(G, (MOD - 1) / len);
            if (invert) w_len = mod_inv(w_len);
            for (int i = 0; i < n; i += len){
                int w = 1;
                for (int j = 0; j < len / 2; j++){
                    int u = a[i+j];
                    int v = mul_mod(w, a[i+j+len/2]);
                    int x = u+v; if (x >= MOD) x -= MOD;
                    int y = u-v; if (y < 0) y += MOD;
                    a[i+j] = x;
                    a[i+j+len/2] = y;
                    w = mul_mod(w, w_len);
                }
            }
        }
        // 3) Divide by n (multiply by inverse) for inverse transform
        if (invert){
            int inv_n = mod_inv(n);
            for (int i = 0; i < n; ++i){
                a[i] = mul_mod(a[i], inv_n);
            }
        }
    }
    // Input : Coefficient vector {a_0, a_1, ...}, {b_0, b_1, ...}
    // Output : Convolution of two coefficient vector
    vector<int> convolution(vector<int> a, vector<int> b) {
        if (a.empty() || b.empty()) return {};
        int S_a = (int)a.size(), S_b = (int)b.size();
        // Make vector size easy to dnc(2^n).
        int n = 1;
        while (n < S_a + S_b - 1){
            n <<= 1;
        }
        a.resize(n); b.resize(n);
        // Normalize to [0, MOD)
        for (int i = 0; i < n; ++i){
            a[i] %= MOD; if (a[i] < 0) a[i] += MOD;
            b[i] %= MOD; if (b[i] < 0) b[i] += MOD;
        }
        // NTT
        iterative_NTT(a, 0); iterative_NTT(b, 0);
        // Pointwise product
        for (int i = 0; i < n; i++) {
            a[i] = mul_mod(a[i], b[i]);
        }
        // INTT
        iterative_NTT(a, 1);
        a.resize(S_a + S_b - 1);
        return a;
    }
};

int main(void){
    fastio

    // Find primes using sieve
    vector<int> sieve(1'000'001, 1);
    sieve[0] = 0; sieve[1] = 0;
    for (int i = 2; i <= 1000; ++i){
        if (sieve[i]){
            for (int j = i*i; j <= 1'000'000; j += i){
                sieve[j] = 0;
            }
        }
    }
    // Product sieve using NTT
    NTT<998'244'353, 3> ntt;
    auto C = ntt.convolution(sieve, sieve);
    int T; cin >> T;
    int N;
    for (int i = 0; i < T; ++i){
        cin >> N;
        cout << (C[N]+1)/2 << endl;
    }
    return 0;
}