halcの競プロ精進ブログ

はるくが競プロしたことを書き留めるなにか

【今日の精進】ABC303Ex - Constrained Tree Degree

毎日格上の問題を倒すやつの49日目です。*1

問題リンク*2

atcoder.jp

問題概要

N 頂点の木のうち、すべての頂点の次数が S に含まれているようなものの個数を 9998244353 で割った余りを求めてね。

解法

2日前にやった次数制限付きケイリーの公式を利用します。

同じように、(N-2)! は後にかけることにします。ので、一旦 \frac{1}{(a_1-1)!(a_2-1)!\dots(a_i-1)!} の部分を考えます。

分解して、\frac{1}{(a_1-1)!}\frac{1}{(a_2-1)!}\dots\frac{1}{(a_N-1)!} とすると、畳み込みができそうな形になります。

S に含まれるところの項だけ \sum_{i=1}^{|S|}\frac{x}{(S_i-1)!}^{S_i} みたいにすればよいのです。

あとはこれを繰り返し二乗法+NTTとかすればよいです。

提出コード

#include<bits/stdc++.h>
#include<atcoder/all>
using namespace std;
using namespace atcoder;
using mint = modint998244353;
int N;
vector<mint> power(vector<mint> &poly,int num){
    if(num==1)return poly;
    vector<mint> half=power(poly,num/2);
    vector<mint> ret=convolution(half,half);
    if(num%2){
        ret=convolution(poly,ret);
    }
    return vector<mint>(ret.begin(),ret.begin()+N);
}

int main(){
    int K;
    cin>>N>>K;
    set<int> S;
    for(int i=0; i<K; i++){
        int S_i;
        cin>>S_i;
        S.insert(S_i);
    }
    vector<mint> poly(N);
    mint fact=1;
    for(int i=0; i<N; i++){
        if(S.count(i+1)){
            poly[i]=fact;
        }
        fact/=i+1;
    }
    mint ans=1;
    for(int i=0; i<N-2; i++){
        ans*=i+1;
    }
    cout<<(ans*power(poly,N)[N-2]).val()<<endl;
}

Submission #47367364 - NS Solutions Corporation Programming Contest 2023(AtCoder Beginner Contest 303)

感想

Pythonの畳み込みは遅い!!!!(憤怒)

*1:もとい、FPSウィークの4日目です。

*2:だいぶ前に解説見ちゃったんだけどゆるして…