結果

問題 No.1066 #いろいろな色 / Red and Blue and more various colors (Easy)
ユーザー Naoto FujiwaraNaoto Fujiwara
提出日時 2023-05-31 15:32:45
言語 C++14
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 52 ms / 2,000 ms
コード長 4,985 bytes
コンパイル時間 1,000 ms
コンパイル使用メモリ 89,536 KB
実行使用メモリ 5,376 KB
最終ジャッジ日時 2024-06-08 21:05:17
合計ジャッジ時間 2,223 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge3
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
5,248 KB
testcase_01 AC 2 ms
5,376 KB
testcase_02 AC 2 ms
5,376 KB
testcase_03 AC 2 ms
5,376 KB
testcase_04 AC 2 ms
5,376 KB
testcase_05 AC 2 ms
5,376 KB
testcase_06 AC 2 ms
5,376 KB
testcase_07 AC 2 ms
5,376 KB
testcase_08 AC 28 ms
5,376 KB
testcase_09 AC 3 ms
5,376 KB
testcase_10 AC 52 ms
5,376 KB
testcase_11 AC 25 ms
5,376 KB
testcase_12 AC 22 ms
5,376 KB
testcase_13 AC 24 ms
5,376 KB
testcase_14 AC 25 ms
5,376 KB
testcase_15 AC 46 ms
5,376 KB
testcase_16 AC 14 ms
5,376 KB
testcase_17 AC 15 ms
5,376 KB
testcase_18 AC 45 ms
5,376 KB
testcase_19 AC 23 ms
5,376 KB
testcase_20 AC 2 ms
5,376 KB
testcase_21 AC 32 ms
5,376 KB
testcase_22 AC 2 ms
5,376 KB
testcase_23 AC 2 ms
5,376 KB
testcase_24 AC 2 ms
5,376 KB
testcase_25 AC 51 ms
5,376 KB
testcase_26 AC 46 ms
5,376 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

/*
AC

(x+A[0]-1)*(x+A[1]-1)*(x+A[2]-1)*...*(x+A[k-2]-1)*(x+A[k-1]-1)
のx^mの係数= ボール1がm個の取り出し方

*/

#include<iostream>
#include<set>
#include<algorithm>
#include<vector>
#include<string>
#include<set>
#include<map>
#include<numeric>
#include<queue>
#include<cmath>
#include<cassert>
using namespace std;
typedef long long ll;
const ll MOD=998244353;


class NTT{

    private:
        long long int mod;
        const long long root=3;
        long long int add(const long long int x, const long long int y)
        {
            return (x + y <  mod ? x + y : x + y -  mod);
        }
        
        long long int sub(const long long int x, const long long int y)
        {
            return (x >= y) ? (x - y) : ( mod- y + x);
        }
        
        long long int mul(const long long int x, const long long int y)
        {
            return x * y %  mod;
        }
        
        long long int mod_pow(long long int x, long long int n)
        {
            long long int res = 1;
            while(n > 0){
                if(n & 1){ res = mul(res, x); }
                x = mul(x, x);
                n >>= 1;
            }
            return res;
        }

        long long int inverse(const long long int x)
        {
            return mod_pow(x,  mod- 2);
        }


        void ntt(std::vector<long long int>& a, const bool rev = false){

                unsigned int i, j, k, l, p, q, r, s;
                const unsigned int size = a.size();
                if(size == 1) return;
                std::vector<long long int> b(size);
                r = rev ? ( mod- 1 - ( mod- 1) / size) : ( mod- 1) / size;
                s = mod_pow(root, r);
                std::vector<long long int> kp(size / 2 + 1, 1);
                for(i = 0; i < size / 2; ++i) kp[i + 1] = mul(kp[i], s);
                for(i = 1, l = size / 2; i < size; i <<= 1, l >>= 1){
                    for(j = 0, r = 0; j < l; ++j, r += i){
                        for(k = 0, s = kp[i * j]; k < i; ++k){
                            p = a[k + r], q = a[k + r + size / 2];
                            b[k + 2 * r] = add(p, q);
                            b[k + 2 * r + i] = mul(sub(p, q), s);
                        }
                    }
                    swap(a, b);
                }
                if(rev){
                    s = inverse(size);
                    for(i = 0; i < size; i++){ a[i] = mul(a[i], s); }
                }
            }

    public:
        NTT(long long int mod_=998244353): mod(mod_){}; 

        //
        //vector<ll> c = ntt(a,b)みたいにして使う
        std::vector<long long int> operator()(const std::vector<long long int>& a, const std::vector<long long int>& b){
            const int size = (int)a.size() + (int)b.size() - 1;
            long long int t = 1;
            while(t < size){ t <<= 1; }

            std::vector<long long int> A(t, 0), B(t, 0);
            
            for(int i = 0; i < (int)a.size(); i++){ A[i] = a[i]; }
            for(int i = 0; i < (int)b.size(); i++){ B[i] = b[i]; }
            
            ntt(A), ntt(B);
            for (int i = 0; i < t; i++){ A[i] = mul(A[i], B[i]); }
            ntt(A, true);
            A.resize(size);
            
            return A;
        }


        //vector<ll> c;
        //  res[j]= ∑(a[i+j]*b[i])(i=[0,b.size()))  j=[0,a.size())

        std::vector<long long int>  filter(const std::vector<long long int>& a,const std::vector<long long int>& b){
            std::vector<long long int> revb=b;
            int lena=a.size(),lenb=b.size();
            
            assert(lena>=lenb);
           

            std::reverse(revb.begin(),revb.end());
            
            std::vector<long long int> filt=this->operator()(a,revb);

            std::vector<long long int> res;
            for(int i=lenb-1;i<lena;i++){
                //[lenb-1,lena)
                res.push_back(filt[i]);
            }
            return res;
            

        }
};


int main(){
    int N,Q;
    cin>>N>>Q;
    vector<ll> A(N),B(Q);
    for(int i=0;i<N;i++){
        cin>>A[i];
      
    }
    for(int i=0;i<Q;i++){
        cin>>B[i];
    }
    NTT ntt;

    for(int i=0;i<N;i++){
        A[i]--;
        A[i]%=MOD;
    }

    //[left,right)
    auto dfs=[&](auto f,int left,int right)->vector<ll>{
        if(right-left==1){
            vector<ll> res(2,1);
            res[0]=A[left];
            //x^1の係数が1, x^0の係数がA[left]
            return res;
        }
        
        int mid=(left+right)/2;

        vector<ll> vl=f(f,left,mid);
        vector<ll> vr=f(f,mid,right);
        
        vector<ll> res=ntt(vl,vr);
        return res;
    };

    vector<ll> ans=dfs(dfs,0,N);

    // ans[i]=色1の玉がi個含まれる場合の数. x^iの係数に対応

    for(int i=0;i<Q;i++){
        cout<<ans[B[i]]<<endl;
    }

    // for(int i=0;i<ans.size();i++){
    //     cout<<"ans["<<i<<"]="<<ans[i]<<endl;
    // }

}
0