Zenn
Closed9

全方位木dpの学習

競技プログラミング
AtCoder
アルゴリズム
動的計画法
全方位木dp
qdot3qdot3

できること

  • すべての頂点を根としたとき、ある種の計算を(N)で実行する。
  • ナイーブにはO(N^2)かかる

手順

  • ある頂点を根として木dp
  • 頂点を１つずつずらして差分を計算

制約

  • 差分計算を効率的に（＝累積的に）行うために、モノイド（結合律・単位元の存在）が対象。
  • ちなみに、単位元は初期化に必要。結合律は、左右からの累積和を結合するときに必要。

https://algo-logic.info/tree-dp/

qdot3qdot3

題材

https://atcoder.jp/contests/dp/tasks/dp_v

N 頂点の木があります。 頂点には 1,2,…,N と番号が振られています。 各 i (1≤i≤N−1) について、i 番目の辺は頂点 x_ i と y_iを 結んでいます。
太郎君は、各頂点を白または黒で塗ることにしました。 このとき、どの黒い頂点からどの黒い頂点へも、黒い頂点のみを辿って到達できるようにします。

正整数 M が与えられます。 各 v (1≤v≤N) について、次の質問に答えてください。

頂点 v が黒であるような頂点の色の組合せは何通りか？ M で割った余りを求めよ。

補足

  • 単にモノクロに塗り分ける方法は 2^N 通りあるので、剰余をとるのは合理的
  • 1≤N≤10^5
qdot3qdot3

入力

use proconio::{input, marker::Usize1};

fn main() {
    input! { n: usize, m: u32 }

    let edge = {
        let mut edge = vec![Vec::new(); n];
        for _ in 1..n {
            input! { x: Usize1, y: Usize1 }
            edge[x].push(y);
            edge[y].push(x);
        }
        edge
    };

    // rerooting
}
qdot3qdot3

考察

どの黒い頂点からどの黒い頂点へも、黒い頂点のみを辿って到達できる

  • 黒い頂点のみを取り出した部分グラフは、木でなければならない。
  • 結局、各ノードについて、それを親とする部分木であるもの数を求めればよい。
  • ノード 0 を根とする部分木の総数は、その子を頂点とする部分木の数を d_i とおくと \prod_i (d_i+1) とかける。

累積的な計算

  • 各ノードについて、子を一列に並べたときに左右から行う

注意事項

  • 順序を考えるのは大変そうなので、可換モノイドに限定して考える。つまり、a\times b=b\times aの性質を利用する。割り算や文字列の結合では成り立たない。
  • モジュラ逆数が存在しない場合があるため、割り算は実質的に使用できない。逆に割り算が使用できる場合、「累積的な計算」で手抜き（＝総積を一度だけとる）ができた。
  • MIN・MAXはいつもさぼれない
  • ADDはいつでもさぼれる
  • ADD・MUL・MAX・MINはどれも可換だから楽できる。
qdot3qdot3

step１：ノード 0 について計算

DFSで実行できる。コンパイル言語でも再帰よりループの方が速いと信じてる。

    // DFS starting at the node 0
    let mut dp = vec![0; n];
    let mut stack = Vec::with_capacity(n);
    stack.push(0);
    let mut is_visited_once = vec![false; n];
    while let Some(&i) = stack.last() {
        if is_visited_once[i] {
            // postorder
            is_visited_once[i] = false;
            stack.pop();

            dp[i] = edge[i]
                .iter()
                .filter(|&&j| !is_visited_once[j])
                .fold(1, |acc, &j| acc * (dp[j] + 1) % m);
        } else {
            // preorder
            is_visited_once[i] = true;
            for &j in edge[i].iter().filter(|&&j| !is_visited_once[j]) {
                stack.push(j)
            }
        }
    }

    println!("0: {}", dp[0]);
qdot3qdot3

step2：差分計算

  • 累積積の再利用を考えるとBFSがよい。メモすると不要なメモリを使用する。
  • 計算の順序が大事
  1. 先ず、累積積を左右から計算する。
  2. 次に、dpを書き換える。
  • dp[i]はノードiが黒であることを仮定しているので、必要なら白の場合を追加する。
    // Rerooting: BFS starting at the node 0
    let mut deque = VecDeque::with_capacity(n);
    deque.push_back((0, 1));
    let mut is_visited = vec![false; n];
    while let Some((i, left)) = deque.pop_front() {
        is_visited[i] = true;

        let children = Vec::from_iter(edge[i].iter().filter(|&&j| !is_visited[j]).copied());
        // println!("children: {:?}", children);

        let mut acc_l = vec![1; children.len()];
        for i in 1..children.len() {
            acc_l[i] = acc_l[i - 1] * (dp[children[i - 1]] + 1) % m
        }
        let mut acc_r = vec![1; children.len()];
        for i in (1..children.len()).rev() {
            acc_r[i - 1] = acc_r[i] * (dp[children[i]] + 1) % m
        }
        // println!("{:?}\n{:?}", acc_l, acc_r);

        for (i, c) in children.into_iter().enumerate() {
            let left = acc_l[i] * acc_r[i] * left + 1;
            dp[c] = (dp[c] * left) % m;
            deque.push_back((c, left))
        }
        // println!("{:?}\n", dp)
    }

    println!("{}", dp.into_iter().join("\n"))
qdot3qdot3

修正

上記step2はオーバーフローすることがある。

-           let left = acc_l[i] * acc_r[i] * left + 1;
+           let left = acc_l[i] * acc_r[i] % m * left % m + 1;
qdot3qdot3

感想

  • 書き出しながら学ぶと理解がはかどる。実際、初めて全方位木DP理解した。
  • 再帰でなくループを使ったからこそ高速かつ省メモリな実装になったと（盲目的に）信じてる。再帰って、それ以前の情報を保持するのにコストかかってそうだし。
  • DFSの帰りがけ順にソートしても解けそう。step1のついでにできて、step2でis_visitedしなくてよいし、累積積もうまく管理できそう。

課題

  • 今回は掛け算（可換な演算）だから簡単だった。可換でない演算についてもできるらしいが...？
このスクラップは5ヶ月前にクローズされました