ながめも

競技プログラミングについて

トップ > AtCoder > AtCoder Beginner Contest 160 F - Distributing Integers

AtCoder Beginner Contest 160 F - Distributing Integers

AtCoder ABC DFS 全方位木DP

F - Distributing Integers

タイトルの上に全方位DPタグがあると思うので、そこから類題を探してもいいと思います。類題が見たくない人は実装の下まで見ないように注意してください。

問題概要

問題へのリンク

見た瞬間に全方位DPが思いつきます。 k = 1の場合さえ求めることができればと思っていたのですが、本番は処理の複雑さにやられてしまいました。

方針

k = 1パート

全方位DPをします。 まず、根が1の根付き木だとして、以下のようなDPの値を計算します。このパートでは、 k = 1の場合のみ答えが求まります。そして、この k = 1の結果から、更新を工夫することで全ての頂点からの場合の数について \mathcal{O}(N)で解くことを考えます。

dp_v := v以下の部分木の場合の数

sub_v := v以下の部分木に含まれる頂点数(dfs0()で求めます。)

とすると、

dp_{cur}は以下のような計算式で求めることができます。

f:id:coonevo:20200329032735j:plain

この説明は、解説放送を聞いていただけるとわかると思います。まずまだ塗られていない点に関して、色を分配して、そこから内部で塗り方をよしなに決めてもらうという方法です。ここまで綺麗に書けると思いませんでした。考えてみれば、子の塗り方は独立なので、別々に考えてかけ算していいです。この考え方前にもやったことあるような気がします。

色々と話がそれましたが、この計算は綺麗なDFS(dfs1)で実装できます。

更新パート

更新ですが、こちらも全方位DPではDFS(dfs2)で実装します。DFSで根から辺を降りながら、逆向きに辺を張り、計算の差分を考えましょう。このパートも説明が難しいですが、繋ぎ換える前にどのような計算をしたのかを考えると自ずと更新式が導かれるのではないかと思います。

以上のような方針により、 \mathcal{O}(N)で解くことができました。

実装

提出へのリンク

int N;
vec<mint> dp, ans;
vvec<int> G;
vec<int> sub;

//cur以下の部分木の頂点数 sub[cur]
int dfs0(int cur = 0, int pre = -1){
    int res = 1;
    for(auto& nv: G[cur]){
        if(nv == pre)continue;
        res += dfs0(nv, cur);
    }
    return sub[cur] = res;
}

//from 0のdpを求める
mint dfs1(int cur = 0, int pre = -1) {
    mint res = C.fact[sub[cur]-1];
    for(auto& nv: G[cur]){
        if(nv == pre)continue;
        res *= C.ifact[sub[nv]];
        res *= (dfs1(nv, cur));
    }
    return dp[cur] = res;
}

void dfs2(int cur = 0, int pre = -1){
    //cerr << pre+1 <<  " -> " << cur+1 << endl;
    //print(dp);
    //print(G[cur]);
    for(auto& nv: G[cur]){
        if(nv == pre)continue;
        mint tmp1 = dp[cur];
        mint tmp2 = dp[nv];

        int tmp_sub1 = sub[cur];
        int tmp_sub2 = sub[nv];

        sub[cur] -= sub[nv];
        sub[nv] =  N;
        //
        dp[cur] *= tmp2.inv(); //dpの値を消す
        dp[cur] *= C.fact[tmp_sub2]; //分母の階乗を消す
        dp[cur] *= C.fact[tmp_sub1-1].inv();// 分子を一旦消す
        dp[cur] *= C.fact[sub[cur]-1];//階上を小さくしてかけ算

        dp[nv] *= dp[cur];// 更新された親のdpの値をかける
        dp[nv] *= C.fact[sub[cur]].inv();//分母に追加
        dp[nv] *= C.fact[tmp_sub2-1].inv();//一旦分子を消す
        dp[nv] *= C.fact[sub[nv]-1];//分子に大きくしてかける
        dfs2(nv, cur);
        //
        dp[cur] = tmp1;
        dp[nv] = tmp2;
        sub[cur] = tmp_sub1;
        sub[nv] = tmp_sub2;
        //cerr << nv+1 <<" -> " << cur+1 << endl;
    }
    ans[cur] = dp[cur];
}
int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(20);

    cin >> N;
    G.resize(N);

    rep(i,N-1){
        int a,b;
        cin >> a >> b;
        a--;b--;
        G[a].push_back(b);
        G[b].push_back(a);
    }
    sub.resize(N);
    dfs0();
    //print(sub);

    dp.assign(N,-1);
    dfs1();
    //print(dp);
        
    ans.resize(N);
    dfs2();
    rep(i,N){
        cout << ans[i] << endl;
    }
}

類題