Rustコトハジメ

プログラミング言語Rustと競プロに関する情報をお届けします。

【Educational DP-V Subtree】全方位木DP

V - Subtree

問題

木が与えられる。初めノードは白に塗られている。

各ノードをルートにした時、ルートから連続して黒に塗る方法は何通りあるか?答えはMで割って答えなさい。

制約: N=105, M=109

方針

各ノードをルートにした時にそこから連続して黒で塗る方法は何通りあるかという問題と見ることが出来ます。

これは、全方位木DPという、知らなくても方針は立つけど、知らないととんでもなく難しくなるDPです。知識なしで実装は結構難しいのではないでしょうか。私は失敗しました。

dp[u][i]をuを親とした時のi番目のノードを黒く塗る通りとします。

dp[u][i]が求まってるとすると、求めるべくは \proddp[u][i]+1になります。なぜならば、各行き先で、そこから黒に塗るか、全部真っ白のままにするか(+1)選べるからです。

ここまではわかる人が多いと思いますが、一体これをどうやって実装するかが難しいです。何をすればいいか考えます。

f:id:akiradeveloper529:20191001214029j:plain

明らかに全ノードに対して新規にDFSするのは間違いです。O(N2)なため。

なので、やりたいことは、ノード0からDFSしたあとに、その情報を使って計算を省きつつ、他のノードをルートにした情報を作りたいということです。これは図を描けばわかるとおり、0->1を逆に1->0にすることと同じです。

これは、ノード0が積xyzと確定したとすると、最初にDPした時に計算したxで割ることでyzを作ることが出来そうですが、Mが素数ではないためこの方針はうまく行きません。(素数ならば出来ます)

ではどうするかというと、割り算を使わないということですが、そのためには毎回掛け算をするということが考えられます。ここでいうと、yzを計算するということです。しかしこれはO(N2)コースですからうまく行きません。

正しい方針は、左からと右からの累積積のようなものを計算して、それを使うと、ノード1を計算した時、ノード2を計算した時、ノード3を計算した時もすべてO(1)で計算することが出来ます。

このアルゴリズムは、ちょうどfoldのようなものなので、きれいに抽象化出来ます。

実装

static mut M: u64 = 0;
fn solve() {
    input!{
        N:usize,m:u64,
        xy:[(usize,usize);N-1]
    }
    unsafe {
        M = m;
    }
    let mut g = vec![vec![];N];
    for (x,y) in xy {
        g[x-1].push(y-1);
        g[y-1].push(x-1);
    }
    let mut reroot: ReRooting<F> = ReRooting::new(g);
    reroot.dfs(0);
    for u in 0..N {
        println!("{}", reroot.solve(u));
    }
}
struct F {}
impl Foldable for F {
    type T = u64;
    fn identity() -> u64 {
        1
    }
    fn fold(acc: &u64, x: &u64) -> u64 {
        unsafe {
            (acc * (x+1)) % M
        }
    }
    fn merge(a: &u64, b: &u64) -> u64 {
        unsafe {
            (a * b) % M
        }
    }
}

ReRootingクラスの実装は、こちらにあります。

GitHub - akiradeveloper/rust-comp-snippets