Rustで動的計画法：Indipendent Set

はじめに

動的計画法を実装してみて、Rustの勉強をやってみる。問題としてはEducational DP Contestという動的計画法の練習を目的としたコンテストのものを使用。

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

AからZまで問題が設定されているが、今回はPのIndipendent Setを解いてみる。グラフ理論での「独立集合」の数を求める問題。ただ、入力が「木」である。つまり閉路がないと制限が入っているので、比較的容易に解ける。

利用したライブラリ

標準入力からデータを読み取る部分はproconioを利用。
https://docs.rs/proconio/latest/proconio/

ベクタから条件にあう要素の位置を得るためにitertoolsに入っているpositionsを使う。
https://docs.rs/itertools/latest/itertools/

完成したコード

長いので、フルサイズはgistに置いている。structとimplでクラスっぽく書いてみた。

https://gist.github.com/attgm/6a63c62b2806bd88e1064f862cfced43

アルゴリズム

ノードを表現するために以下のような構造体を用意している。隣接ノードの集合neighborsと、親ノードが黒の時と白の時とで、自分とその子ノードが取りうる組み合わせの数valueからなる。

independent_set.rs

struct Node {
    neighbors: Vec<usize>,
    value: (usize, usize)
}

入力としては木なので、あるノード($n_k$)の隣接ノードは1つの親ノード($p$)と複数の子ノード($C = \{ c_i \}$)に分けられる。この時、親が黒の場合、隣接する自ノードは必ず白なので、

$$B(n_k) = \prod_{\forall c_i} W(c_i)$$

親が白の場合、自ノードは白と黒の2通りが考えられるので、

$$W(n_k) = \prod_{\forall c_i} W(c_i) + \prod_{\forall c_i} B(c_i)$$

で求められる。ただし、ノード$x$の親ノードが黒の時と白の時とで、ノード$x$とその子ノードが取りうる組み合わせの数をそれぞれ$B(x), W(x)$とする。つまり、valueは$(B(x), W(x))$である。

後は全ての子ノードの$W(c_i), B(c_i)$が決まった段階で、自身の$W(n_k), B(n_k)$を求めていき、ルートノード$r$は白黒どちらでも良いので、$W(r)$が求める解答となる。

処理の流れ

リーフノードは、子ノードを持っていないので最初に計算が可能となる。初期値として隣接ノードが1つしかないリーフノードをまずqueueに入れる。

independent_set.rs

let mut queue: Vec<_> = nodes.iter()
    .positions(|node| node.neighbors.len() <= 1)
    .collect();

nodesのインデックスがノードのIDになっているので、positionsを使っている。汎用性を持たせるならNodeにIDを入れておくべきかも知れない。条件が==ではなく<=になっているのは、ノードが1つだけの入力に対応するため。

実際の処理は以下のような感じ。queueはLongest Path でやったようにVecDequeを使ってFIFOを作っても良いのだが、queueを切り替えていく感じで実装した。

independent_set.rs

let answer = 'main_loop: loop {
    let mut next_queue = Vec::new();
    for i in &queue {
        let index = *i;
        if let Some(parent) = nodes[index].set_parent_node() {
            let value = nodes[index].value;
            if nodes[parent].set_child_value(index, value) {
                next_queue.push(parent);
            }
        } else {
            break 'main_loop nodes[index].value.1;
        }
    }
    queue = next_queue;
    assert!(queue.len() > 0);
};

queueに入っているノードは「全ての子ノードが計算済みのノード」なので、set_parent_node()でvalueを計算し、親ノードparentを取得する。親ノードparentに子ノードiの終了を知らせ、parentのすべての子ノードの計算が完了していたら、next_queueにparentを挿入する。

すべての子ノードの計算が終了したが親ノードが存在しない場合、このノードがルートノードなので、
結果を返して終了する。ここで単にbreakだけだとforのループからしか抜けないので、loopにmain_loopという名前をつけて外側のループまで抜けるようにしている。さらにRustのloopは返り値を得ることができるので、breakで渡すようにしている。ループが値を取れるというのは、あまり見ない気はするが、今回の処理にはピッタリはまっている。変数のスコープなどを考えるとちょっと分かりにくい記述になりがちなので、スッキリ書けるのはいい。

親子関係

入力は「木」であることは保証されているが、ルートノードの指定はない。木の場合、任意のノードをルートノードにできるので実はどれがルートノードでも構わない。なので、親子関係についても「一番最後まで計算が終わらなかったノードを親とする」という方法で決定している。

具体的には以下の関数になる。計算が終わった子ノードidを隣接ノードneighborから取り除いていって、最後の1個になった時点で「全ての子ノードの計算が終わった」と返す。このノードが参照される前に、最後の隣接ノードの計算が完了し、隣接ノードが0になった場合は、このノードはルートノードになる。それでうまくいくの？という気はするが、閉路がなく、すべてのノードが連結していれば、これでうまくいく。

independent_set.rs

fn set_child_value(&mut self, id: usize, value:(usize, usize)) -> bool {
    if let Some(index) = self.neighbors.iter().position(|x| *x == id) {
        self.neighbors.swap_remove(index);
        self.value = ((self.value.0 * value.0) % MOD, (self.value.1 * value.1) % MOD);
    }
    self.neighbors.len() == 1
}

あまり良くない使い回し

あまり良くないとは思うが、今回Node構造体のvalueについて2つの意味を持たせている。
set_child_valueの中では、途中計算結果を保持するためvalueは以下を表している。

$$\left(\prod_{\forall c_i} B(c_i) , \prod_{\forall c_i} W(c_i)\right)$$

set_parent_nodeで途中経過結果を使って以下に置き換えて、value本来の定義に変更している。

$$\left(B(n_k) = \prod_{\forall c_i} W(c_i), W(n_k) = \prod_{\forall c_i} B(c_i) + \prod_{\forall c_i} W(c_i)\right)$$

set_parent_nodeの呼び出し後でないとvalueが参照されないので問題ないのだが、可読性という意味ではあまり良くない。とはいいつつ、途中経過を示すためのメンバ変数を追加するのも冗長な気はする。

計算処理時間の測定

問題の生成

ノードの数$N$に対して、$N-1$本のエッジからなるグラフを生成する必要があるので、$N-1$までのノードから1本ずつエッジを張ることで生成する。さらに自分のノードIDより大きいノードとしかエッジを張らないようにすることで、閉路が生成されない。この時の確率分布に差をつけるともっといいのかも知れないが、一様分布で生成した。

今回のコードではあまり関係ないが、ノードやエッジのでてくる順番で差がでてこないように、ノードIDをランダムで割り当てるようにし、さらに表示する順番にもランダムを入れた。

gen_indipendent_set.py

#!/user/bin/env python
import random
import sys

n = int(sys.argv[1])

edge = [random.randrange(i + 1, n) for i in range(n - 1)]
node_map = list(range(1, n + 1))
random.shuffle(node_map)
order_map = list(range(0, n - 1))
random.shuffle(order_map)

print(n)
for i in order_map:
    print(str(node_map[i]) + " " + str(node_map[edge[i]]))

測定結果

両対数グラフで書いているが、綺麗に$O(N)$のグラフになっている。$N=10^5$でも10ms以下なので速度的にも問題ないと思われる。

関連記事

Rustで動的計画法の実装:
🐸Frog | 🌴Vacation | 🎒Knapsack | 🐍LCS | 🚶‍♂️Longest Path | 🕸️Grid | 💰Coins | 🍣Sushi | 🪨Stones | 📐dequeue | 🍬Candies | 🫥Slimes | 💑Matching | 🌲Indipendent Set | 🌻Flowers | 👣Walk | 🖥️Digit Sum | 🎰Permutation | 🐰Grouping | 🌿Subtree | ⏱️Intervals | 🗼Tower | 🐸Frog3