結果

問題 No.1614 Majority Painting on Tree
ユーザー momoyuu
提出日時 2024-08-09 21:02:52
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 1,487 ms / 5,000 ms
コード長 1,673 bytes
コンパイル時間 1,444 ms
コンパイル使用メモリ 115,032 KB
実行使用メモリ 14,592 KB
最終ジャッジ日時 2024-08-09 21:03:20
合計ジャッジ時間 22,715 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 4
other AC * 45
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include<iostream>
#include<vector>
#include<algorithm>

using namespace std;
using ll = long long;

#include<atcoder/modint>
using mint = atcoder::modint998244353;

int main(){
    cin.tie(nullptr);
    ios::sync_with_stdio(false);

    int n,c;
    cin>>n>>c;
    vector<vector<int>> g(n);
    for(int i = 0;i<n-1;i++){
        int u,v;
        cin>>u>>v;
        u--;v--;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    int nni = 1;
    vector<mint> nipow(n+1,1),njpow(n+1,1);
    mint inv;
    vector<mint> fac(n+1,0),ifac(n+1,0);
    fac[0] = 1;
    for(int i = 1;i<=n;i++) fac[i] = fac[i-1] * i;
    ifac[n] = fac[n].inv();
    for(int i =  n-1;i>=0;i--) ifac[i] = ifac[i+1] * (i+1);
    auto dfs = [&](auto dfs,int ni,int p) -> mint {
        int all = g[ni].size();
        mint sum = 1;
        for(auto&to:g[ni]) if(to!=p) {
            mint tmp = dfs(dfs,to,ni);
            sum *= tmp;
        }
        mint ans = nipow[all] * sum;
        for(int i = 1;i<all;i++){
            if(2*i<=all) continue;
            mint tmp = fac[all] * ifac[i] * ifac[all-i];
            tmp *= nni;
            tmp *= njpow[all-i];
            tmp *= sum;
            ans -= tmp;
        }
        return ans * inv;
    };

    mint ans = 0;
    for(int i = 1;i<=c;i++){
        nni = i;
        inv = mint(nni).inv();
        for(int j = 1;j<=n;j++){
            nipow[j] = nipow[j-1] * nni;
            njpow[j] = njpow[j-1] * (nni-1);
        }
        auto use = dfs(dfs,0,-1);
        use *= nni;
        if((c-i)%2==0) ans += use * fac[c] * ifac[i] * ifac[c-i];
        else ans -= use * fac[c] * ifac[i] * ifac[c-i];
    }
    cout<<ans.val()<<endl;
}
0