2022牛客寒假算法基础集训营6 牛妹的数学难题 排列组合+逆元

#include<bits/stdc++.h>
using namespace std;
#define int long long 
const int N=1e7+10;
const int mod=998244353;
int n,k,a[N];
int b[5];
int ksm(int a,int b){
    int ans=1;
    while(b){
        if(b%2) ans=ans*a%mod;
        a=a*a%mod;
        b=b>>1; 
    }
    return ans;
}
int p[N],p2[N],invp[N];
int inv(int x){
    return ksm(x,mod-2);
}
int C(int n,int m){
    if(n<m) return 0;
    // cout<<n<<" "<<m<<endl;
    // cout<<(p[n]*inv(p[m])%mod)*inv(p[n-m])%mod<<endl;
    return (p[n]*invp[m]%mod)*invp[n-m]%mod;
}
int32_t main(){
    p[0]=p[1]=p2[0]=1;
    for(int i=1;i<N;i++){
        p[i]=p[i-1]*i;p[i]%=mod;
        p2[i]=p2[i-1]*2;p2[i]%=mod;
    }
    invp[N-1]=ksm(p[N-1],mod-2);
    for(int i=N-2;i>=0;i--){
        invp[i]=invp[i+1]*(i+1)%mod;
    }
    scanf("%lld%lld",&n,&k);
    for(int i=1;i<=n;i++){
        scanf("%lld",&a[i]);
        b[a[i]]++;
    }
    int c1=b[1],c2=b[2];
    int ans=0;
    for(int i=0;i<=k;i++){
        if(c1<i||c2<k-i) continue;
        ans=(ans+C(c1,i)*C(c2,k-i)%mod*p2[k-i])%mod;
    }
    cout<<ans<<endl;
    
}