概要

人によって実装が大きく分かれることに定評のある全方位木DPですが、私のスタイルの内容と、その背景をご紹介します。なおプログラミング言語特有の設計のお話は、好みが分かれるところだと思いますためこの記事では敢えて省いております。

内容

次のコードでできます。演算が可逆でない場合にも、非可換な場合（そんな場合があるのかは疑問ですが）にも対応した万能スタイルでなんとたったの 21 行、うれしいですね。累積和の配列等、無駄な動的メモリ確保がないのも推しポイントです。

let mut lower = vec![Value::identity(); n];
let mut branch = vec![Value::identity(); n];
for &i in sorted.iter().rev() {
    for &j in &g[i] {
        lower[i] = lower[i].mul(&branch[j]);
    }
    branch[i] = lower[i].up();
}
let mut upper = vec![Value::identity(); n];
for &i in &sorted {
    let mut suffix = upper[i].clone();
    for &j in g[i].iter().rev() {
        upper[j] = suffix.clone();
        suffix = branch[j].mul(&suffix);
    }
    let mut prefix = Value::identity();
    for &j in &g[i] {
        upper[j] = upper[j].mul(&prefix).up();
        prefix = prefix.mul(&branch[j]);
    }
}

必要なメソッド

実装する必要のあるメソッドは次の 3 つです。トレイトにするとかそういう流派がわかれそうな設計のお話は今回はしません。今回は深入りしませんが、節点重みが必要な場合は up() の引数に追加するなどするとよいと思います。辺重みは節点重みに変換するなどしましょう。

impl Value {
    fn up(&self) -> Self;
    fn identity() -> Self;
    fn mul(&self, other: &Self) -> Self;
}

3 つの変数の意味

lower, branch, upper は次のような根付き木リストに対応する値であると考えます。（変数名微妙だと思っているため、名案募集中です。）

	対応するオブジェクト
`lower[i]`	$i$ の子に対応する部分木のリスト
`branch[i]`	$i$ を根とする部分木のみからなるシングルトン
`upper[i]`	$i$ の親を根とする上向きの木があればそれのシングルトン、なければ空リスト

3 つのメソッドの意味

このとき up(), identity(), mul() は次のような演算に対応します。

	型（$T$ は根付き木型）	対応するオブジェクト
`up()`	$[T] \times V \to [T]$	リスト内のすべての根付き木を、新しい接点で join する
`identity()`	$* \to [T]$	空リストを返す
`mul()`	$[T] \times [T] \to [T]$	リストの結合（concatenation）

もともと定義されていないといけない変数

しれっと書いてありますが、こういうものを構築しておく必要があります。

	型	意味
`g`	`Vec<Vec<usize>>`	子リスト（親が消されていることに注意）
`sorted`	`Vec<usize>`	節点の順列であって、親が子よりも先に出てくるもの

初期変数の構築

例えばこういう感じでやるとよいです。毎回同じコードなので爆速タイピングできるようにしておきましょう。特に sorted, parent, stack の 3 つ組はしょっちゅう使うのでセットで暗記しています。

proconio! {
    n: usize,
    edges: [(Usize1, Usize1); n],
}

let mut g = vec![Vec::new(); n];
for &(i, j) in &edges {
    g[i].push(j);
    g[j].push(i);
}

let mut sorted = Vec::new();
let mut parent = vec![usize::MAX; n];
let mut stack = vec![0];
while let Some(i) = stack.pop() {
    sorted.push(i);
    g[i].retain(|&j| parent[i] != j);
    for &j in &g[i] {
        stack.push(j);
        parent[j] = i;
    }
}

具体的な使い方

これも記事の主題ではないためあまり深くは言及しないのですが、たいていの場合は次の 2 通りのどちらかでうまくいくと思います。

upper[i].mul(&lower[i]) を使う
upper[i], branch[j] for j in g[i] を使う

lower と branch は比較的簡単に相互変換できる（特に lower から branch は up() するだけ）ことに着目すると、特に使い方 2 のときは lower を定義せず branch を使いまわすことで若干行数が減ると思いますが、このあたりもあまり深追いしないことにします。繰り返しになりますがこの記事はライブラリ化するときの設計を説明するものではありませんため、そのあたりは適宜ご自由にです。

背景

実装こそ大きく異なるものの、木 DP の考え方自体は抽象化全方位木DPのライブラリとドキュメント | 東京工業大学デジタル創作同好会traP にほとんど共感します。https://x.com/shogo3142/status/1806279294840459689 で言及されている通り、私も関数は $3$ つで考えるのが良いと思います。しかし上記の説明では関数は $2$ つ（identity() はノーカンなので、mul() と up()）しか使っていませんね。そのあたりの対応をご説明しましょう。

DP の遷移が次のように表されるとしましょう。（ただし $j \lessdot i$ は $j$ が $i$ の子であるという意味の独自の記号）

$$ \mathtt{dp}(i) = g\left( \prod _ { j \lessdot i } f(\mathtt{dp}(j) ) \right) $$

$\mathtt{dp}(i)$ は $i$ を根とする部分木 $T _ i ^ { ▼ }$（これも独自の記号）に対応する値なので、$\phi(T _ i ^ { ▼ })$ と表しておきましょう。さらに $i$ の親を根とする上向きの部分木を $f(T _ i ^ { △ } )$ と表します。このとき lower, branch, upper は次のような値を管理しています。なお up() は $f \circ g$ に対応しています。グラフの言葉でいうと $f$ は根付き木を受け取ってそのシングルトンを返す関数、$g$ は根付き木のリストを受け取って新しい節点で join した根付き木を返す関数に対応します。根付き木型と根付き木リスト型が登場していることに注意です。

$$ \mathtt{lower}(i) = \prod _ { j \lessdot i } f(\phi(T _ j ^ { ▼ })) \\ \mathtt{branch}(i) = f(\phi(T _ i ^ { ▼ })) \\ \mathtt{upper}(i) = f(\phi(T _ i ^ { △ })) $$

従って、木の根を $i$ に取り換え（reroot）てできる木を $T _ i$ とおくと、全方位木 DP の値は次のように求まります。

$$ \mathtt{dp}(T _ i) = g \left( \mathtt{lower}(i) \cdot \mathtt{upper}(i) \right) $$

しかしながらこの理論の通りにそのまま実装すると、$f$ と $g$ を分けて実装しないと行けなかったりなどなど微妙な点が多かったため、使いやすさ・実装しやすさ・無駄のなさに寄せて大きく改変した結果が、冒頭のコードです。

ブログ名

競技プログラミングやお歌のお話をする高菜です。

【高菜式】Easiest 全方位木 DP

概要

内容