[C++] 28866번 Морти покупает продукты

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

---

 

26/02/18

 

 

자주 나왔던 단골 유형의 FFT문제이다.

---

 

문제 접근 방식:

 

 

$N$개의 수들 중 중복을 포함하여 $K$개를 선택했을 때의 가능한 경우의 수를 세는 문제라고 요약할 수 있다.

이전에 너무 많이 했던 FFT문제 유형이다. 특히, 보석 가게 문제와 거의 유사하다고 볼 수 있다.

실제로 필요한 다항식의 차수는 $50\ 000$까지 밖에 안되므로, 거듭제곱을 하며 중간 중간 크기를 resize해줘도 상관 없다.

또한, 여기서 $786\ 433$이라는 특이한 수가 쓰이는데, 이 수는 $3\times 2^{18}+1$로 NTT-friendly한 소수이다. 원시근은 $10$을 사용하면 된다.

이 점만 유의하여 구현하고, 나온 결과를 누적합을 한 뒤 각 쿼리를 처리하면 충분하다.

---

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

// 28866번 Морти покупает продукты
// NTT
#include <iostream>
#include <vector>

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

/*
MOD = 786'433 = 3*2^18 + 1 / G = 10 / Works for length up to 2^18 = 262'144
*/
template<int MOD, int G>
struct NTT {
private :
    // HELPER FUNCTION
    inline int mul_mod(int a, int b) const {return (long long)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);}
    // ---- precompute cache (per template instantiation) ----
    static inline vector<int> rev;
    static inline vector<int> roots_fwd, roots_inv;
    static inline int prepared_n = 0; // current prepared size (power of two)
    void precompute(int n) {
        if (prepared_n == n) return;
        prepared_n = n;
        rev.assign(n, 0);
        int lg = 0;
        while ((1 << lg) < n) ++lg;
        for (int i = 0; i < n; i++) {
            rev[i] = 0;
            for (int b = 0; b < lg; b++){
                if (i & (1 << b)) rev[i] |= 1 << (lg - 1 - b);
            }
        }
        roots_fwd.assign(n, 1); roots_inv.assign(n, 1);
        for (int len = 1; len < n; len <<= 1) {
            int wlen_fwd = pow_mod(G, (MOD - 1) / (2 * len));
            int wlen_inv = mod_inv(wlen_fwd);
            int w_f = 1; int w_i = 1;
            for (int j = 0; j < len; j++) {
                roots_fwd[len + j] = w_f;
                roots_inv[len + j] = w_i;
                w_f = mul_mod(w_f, wlen_fwd);
                w_i = mul_mod(w_i, wlen_inv);
            }
        }
    }
public : 
    void iterative_NTT(vector<int>& a, int invert){
        int n = (int)a.size();
        precompute(n);
        // 1) bit-reversal permutation (use rev[])
        for (int i = 1; i < n; i++){
            int j = rev[i];
            if (i < j) swap(a[i], a[j]);
        }
        // 2) butterflies by length = 2, 4, 8, ...
        // use roots_fwd/roots_inv precomputed
        const auto& roots = invert ? roots_inv : roots_fwd;
        for (int len = 2; len <= n; len <<= 1){
            int half = len >> 1;
            for (int i = 0; i < n; i += len){
                for (int j = 0; j < len / 2; j++){
                    int u = a[i+j];
                    int v = mul_mod(roots[half+j], a[i+j+half]);
                    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+half] = y;
                }
            }
        }
        // 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;
    }
    // Input : Coefficient vector A = {a_0, a_1, ...} and Power
    // Output : Coefficient vector of A^P
    vector<int> power(vector<int> a, int e){
        if (a.empty()) return {};
        if (e == 0) return {1};
        int S_a = (int)a.size();
        // Make vector size easy to dnc(2^n).
        int n = 1;
        while (n < (S_a - 1)*e + 1){
            n <<= 1;
        }
        a.resize(n);
        // Normalize to [0, MOD)
        for (int i = 0; i < n; ++i){
            a[i] %= MOD; if (a[i] < 0) a[i] += MOD;
        }
        // NTT
        iterative_NTT(a, 0);
        // Pointwise product
        for (int i = 0; i < n; i++) {
            a[i] = pow_mod(a[i], e);
        }
        // INTT
        iterative_NTT(a, 1);
        a.resize((S_a - 1)*e + 1);
        return a;
    }
};

vector<int> poly_pow(vector<int> poly, int K){
    NTT<786'433, 10> ntt;
    vector<int> ret = {1};
    while (K > 0){
        if (K & 1){
            ret = ntt.convolution(ret, poly);
            ret.resize(50001);
        }
        poly = ntt.convolution(poly, poly);
        poly.resize(50001);
        K >>= 1;
    }
    return ret;
}

int main(void){
    fastio
    
    int N, K, Q; cin >> N >> K >> Q;
    vector<int> poly(50001, 0);
    int c;
    for (int i = 0; i < N; ++i){
        cin >> c;
        ++poly[c];
    }
    auto ret_poly = poly_pow(poly, K);
    for (int i = 1; i < 50001; ++i){
        ret_poly[i] += ret_poly[i-1];
        ret_poly[i] %= 786'433;
    }
    int l, r;
    for (int i = 0; i < Q; ++i){
        cin >> l >> r;
        int query_ans = ret_poly[r] - ret_poly[l-1];
        if (query_ans < 0) query_ans += 786'433;
        cout << query_ans << endl;
    }
    return 0;
}