halcの競プロ精進ブログ

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

【今日の精進】AGC005F - Many Easy Problems

毎日格上の問題を倒すやつの90日目です。

問題リンク

atcoder.jp

問題概要

以下のような問題を考えるよ。

整数 N,KN 頂点の木がある。

頂点を K 個選ぶ方法は _N\text{C}_K 通りあるが、それぞれの選び方に対して、「選んだ頂点すべてを含む最小の全域木の頂点数」をスコアとする。

全ての選び方に対するスコアの合計を 924844033 で割った余りを求めよ。

この問題を一瞬で解いた高橋くんは、K=1,2\dots,N のすべての場合について高速に答えを求める方法を考えることにしたよ。

実際に求めてね。

解法

まず、K が固定されたものを考えます。しかし、これも少し難しいです。

直接求めるのは難しいので、余事象を考えます。というか、「全域木に含まれない頂点数」をスコアにします。これを元に戻すのは容易です。

スコアの対象を変えた後の問題で、各頂点のスコアへの寄与を考えます。これは、その頂点を取り除いたときにできる部分木で、そのうち1つにのみ選んだ頂点が集中しているときに加算されます。

よって、木DPして各頂点ごとに求めれば良いです。後のために、部分木の大きさをすべてメモしておき、頂点数 i の部分木の個数が a_i であるような整数列 a を作っておきます。

K を固定した時の(スコアを言い換えた)答えを改めて考えると、\sum_{i=0}^{N}a_i\times_i\text{C}_K となります。二項定理より、これは [x^K]\sum_{i=0}^{N}a_i\times(x+1)^i となります。

ここまで来れば後一息で、f(x)=\sum_{i=0}^{N}a_i\times x^i としたときの f(x+1) を求めたいので、これは「Polynomial Taylor Shift」に学ぶ畳み込み芸 - けんちょんの競プロ精進記録をそのまま適用できます。

提出コード

#include<bits/stdc++.h>
#include<atcoder/all>
using namespace std;
using namespace atcoder;
using mint = static_modint<924844033>;
using ll = long long;
#define rep(i,n) for(int i=0; i<n; i++)

int dfs(int pos,int bef,int &N,vector<vector<int>> &tree,vector<mint> &dp){
    int ret=1;
    for(int i:tree[pos]){
        if(bef==i)continue;
        int now=dfs(i,pos,N,tree,dp);
        dp[now]++;
        ret+=now;
    }
    dp[N-ret]++;
    return ret;
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int N;
    cin>>N;
    vector<vector<int>> tree(N);
    rep(i,N-1){
        int u,v;
        cin>>u>>v;
        tree[u-1].push_back(v-1);
        tree[v-1].push_back(u-1);
    }
    vector<mint> dp(N+1,0),fact(N+1,1),rev(N+1,1);
    dfs(0,-1,N,tree,dp);
    rep(i,N){
        fact[i+1]=fact[i]*(i+1);
        rev[i+1]=rev[i]/(i+1);
    }
    rep(i,N+1){
        dp[i]*=fact[i];
    }
    vector<mint> pol(N+1);
    rep(i,N+1){
        pol[i]=rev[N-i];
    }
    vector<mint> shift=convolution(pol,dp);
    rep(i,N+1){
        if(i==0)continue;
        cout<<(fact[N]*rev[i]*rev[N-i]*N-shift[i+N]*rev[i]).val()<<"\n";
    }
}

Submission #48649270 - AtCoder Grand Contest 005

感想

最後の最後、辺を有向にしてめっちゃ落としたのは私です