[C++] 13575번 보석 가게

 

 

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

---

 

26/02/10

 

 

FFT를 사용하지 않고도 냅색만으로 해결할 수도 있지만, 나는 그렇게 해결하지 못했기 때문에 FFT를 사용한 풀이만을 소개하겠다.

---

 

문제 접근 방식:

 

 

결국 이것도 Golf Bot문제의 연장선인데, 우리는 보석을 $K$개 가져갈 수 있으므로 $\sum x^{a_{i}}$를 $K$번 거듭제곱하면 된다.

근데 무작정 $K$번 곱한다면 시간초과가 난다.

우리가 Convolution을 구현했을 때를 생각해보면, resize가 필수적이었다.

만약 두 다항식의 길이 차이가 크다면, 길이가 짧은 다항식의 resize 양 또한 커지게 되고, 결국 FFT변환할 때의 시간 또한 오래 걸릴 것이다.

결국 두 다항식의 길이 차이가 "적도록" 다항식을 서로 곱해야 한다.

우리는 분할 정복을 이용한 거듭제곱으로 다항식을 곱할 것이다. 그러면 Convolution함수를 $\log$번 호출하게 되며, 나이브하게 $K$번을 곱하는 것보다 빠르게 구할 수 있다.

---

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

// 13575번 보석 가게
// FFT
/*
접근 방법:
길이 1'001짜리 다항식을 K번 곱한다면?
*/
#include <iostream>
#include <vector>

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

template<int MOD, int G>
struct NTT{
// HELPER FUNCTION
private:
    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 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) butterfly
        for (int len = 2; len <= n; len <<= 1){
            int wlen = pow_mod(G, (MOD-1)/len);
            if (invert) wlen = mod_inv(wlen);
            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, wlen);
                }
            }
        }
        // 3) for invert
        if (invert){
            int inv = mod_inv(n);
            for (int i = 0; i < n; ++i){
                a[i] = mul_mod(a[i], inv);
            }
        }
    }
    vector<int> convolution(vector<int> a, vector<int> b){
        int S_a = (int)a.size(), S_b = (int)b.size();
        int n = 1;
        while (n < S_a+S_b-1){
            n <<= 1;
        }
        a.resize(n); b.resize(n);
        ntt(a, 0); ntt(b, 0);
        for (int i = 0; i < n; ++i) a[i] = mul_mod(a[i], b[i]);
        ntt(a, 1);
        a.resize(S_a+S_b-1);
        return a;
    }
};

vector<int> pow_poly(vector<int> a, int k){
    NTT<998'244'353, 3> ntt;
    vector<int> ret = {1};
    while(k > 0){
        if (k&1) ret=ntt.convolution(ret,a);
        a=ntt.convolution(a, a);
        k>>=1;
    }
    return ret;
}

int main(void){
    fastio

    int N, K; cin >> N >> K;
    vector<int> a(1001, 0);
    int j;
    for (int i = 0; i < N; ++i){
        cin >> j;
        a[j] = 1;
    }
    vector<int> c = pow_poly(a, K);
    for (int i = 0; i < c.size(); ++i){
        if (c[i] != 0) cout << i << ' ';
    }
    return 0;
}