[Rust]petgraphでDFSを実践 (ABC201-E)

4月からRustにハマっています。
AtCoderで使えるクレートにpetgraphがあります。日本語の参考資料に乏しいので、実践例を紹介します。
執筆時点でのバージョンは "0.5.0" です。

問題は、[AtCoder ABC201E](https://atcoder.jp/contests/abc201/tasks/abc201_e) です。
アルゴリズムそのものについては公式サイトを見てください。

https://atcoder.jp/contests/abc201/submissions/22644408

---
# useなど

```rust:
const MOD: usize = 1_000_000_007;

// グラフのコンストラクタ等の基本機能
use petgraph::prelude::*;

// 深さ優先探索を扱う構造体。既に発見した頂点の情報をビットセットの形でself.discoveredとして持つ。
use petgraph::visit::Dfs;

// 頂点のindex(通し番号)から頂点を表す構造体NodeIndexを返す関数
use petgraph::graph::node_index;

// is_visited関数をdiscoveredに対して通すために必要
use petgraph::visit::VisitMap;

// 頂点のidを1-indexedから0-indexedにする
use proconio::{input, marker::Usize1};
```
---
# 入力を読み込んでグラフを構成

```rust:
fn main() {
	input!{ n: usize, uvw: [(Usize1, Usize1, u64); n-1], };
	let mut g = UnGraph::<u64, u64, usize>::from_edges(uvw);
```
辺の数はnではなくn-1なので、`input!{ uvw: [(Usize1, Usize1, u64)] };` だとエラーが出ます。
1-indexedのままコンストラクトすると辺を持たない頂点'0'が出来てしまうので、`Usize1`を利用して0-indexedに自動変換しています。
`UnGraph`は無向グラフを表します。有向グラフは`DiGraph`です。
型指定`<u64, u64, usize>`の一つ目と二つ目の`u64`は、それぞれ頂点と辺の重みです。三つ目は頂点と辺のindex(通し番号)の型です。省略するとデフォルトの`u32`になります。今回`Usize1`を使用したので`usize`を指定しています。
辺情報が(始点index, 終点index, 辺の重み)のタプルの配列(あるいはその他のイテラブルなコレクション)で与えられると、グラフの構成は`from_edges`メソッドで一瞬で済みます。素敵です。

今回、$dist(0, i)$ の記録先として頂点 *i* の重みを割り当てています。普通の配列を用意すればよいのですが、デモンストレーションです。頂点の重み指定は`from_edges`メソッドの引数にはないため、`usize`のデフォルト値の`0usize`が入ります。
重みが不要な場合は、型を`()`で指定します。
辺の重みが不要な場合、`from_edges`の引数は(始点index、終点index)のタプルの配列です。

indexをu32で扱うなら、

```diff:
-	input!{ n: usize, uvw: [(Usize1, Usize1, u64); n-1], };
+	input!{ n: usize, mut uvw: [(u32, u32, u64); n-1], };
+	uvw.iter_mut().for_each((|(ru, rv, _)| { *ru -= 1; *rv -= 1; });
-	let mut g = UnGraph::<u64, u64, usize>::from_edges(uvw);
+	let mut g = UnGraph::<u64, u64>::from_edges(uvw);	
```
などとします。

---
# DFS

```rust
	let mut dfs = Dfs::new(&g, node_index(0));
	while let Some(nx) = dfs.next(&g) {
		let mut nes = g.neighbors(nx).detach();
		while let Some((edge, ny)) = nes.next(&g) {
			if dfs.discovered.is_visited(&ny) { continue; }
			g[ny] = g[nx] ^ g[edge];
		}
	}
```
`Dfs::new(&g, node_index(0))`は、index0の頂点を始点としてグラフg上でDFSを行うインスタンスを返します。`node_index(i)`は頂点であることを示すラッパーを付加しますが、どのグラフについて言及しているかの情報は持っていません。
`dfs.next(&g)`は呼び出すたびにDFSで次の頂点を探索して`Some(NodeIndex)`の形で返し、到達可能な頂点がなくなると`None`を返します。
`g.neighbors(nx)`は、グラフg上でnxと隣接する頂点を返すイテレータとして機能する`Neighbors`を返します。
`.detach()`は、隣接情報を借用を使わずに提供するメソッドで、`Neighbors`に適用して`WalkNeighbors`を返します。これを使用するとグラフを改変（辺や頂点の追加や削除）しながら探索することが可能になります。今回は探索中に改変は行いませんが、`WalkNeighbors`は`.next(&g)`によって隣接する辺と頂点を`(EdgeIndex, NodeIndex)`のタプルで返すので採用しています。通常のイテレータのように`.next()`を適用すると、コンパイルエラーになります。
`discovered`は既に発見した頂点の情報で、型はビットセット`FixedBitSet`です。
`is_visited(&ny)`は、nyを既に訪問したかどうかを返すbool関数です。グラフが木の場合、隣接頂点がDFSの始点側なら`true`、反対側なら`false`になります。
`g[ny]`、`g[nx]`は頂点の重み、`g[edge]`は辺の重みを表します。構文が同じですが引数が`NodeIndex`か`EdgeIndex`かで区別され適切に処理されます。頂点の重みを更新しているのでグラフはミュータブルである必要があり、グラフのインスタンスを作成する際に`let mut g = ...`としています。

`g[_]`の記法を用いない場合、

```diff:
-			g[ny] = g[nx] ^ g[edge];
+			let mut nyw = g.node_weight_mut(ny);
+			*nyw = g.node_weight(nx) ^ g.edge_weight(edge);
```
などとなります。

---
# 位ごとにビットが立っている数をカウント

```rust:			
	let mut cnt = [0; 60];
	for node in g.node_indices() {
		for i in 0..60 {
			if (g[node] >> i) & 1 == 1 { cnt[i] += 1; }
		}
	}
```
`g.node_indice()`は、gの全ての頂点を`NodeIndex`の型で次々返すイテレータを返します。
`g[node]`で頂点の重みにアクセスします。

# 各位ごとに計算して集計、表示

```rust:
	let mut ans = 0;
	for (i, &x) in cnt.iter().enumerate() {
		ans = ( ans + (1<<i) % MOD * x * (n-x) ) % MOD;
	}

	println!("{}", ans % MOD);
}
```
`iter()`は各要素に対する参照を返すので、返り値を`&x`にマッチさせると`x`が値のCopyになります。`enumerate()`はインデックスを(`&usize`でなく)`usize`で付加するので、全体の型は`(usize, &usize)`です。

---

実戦中は`WalkNeighbors`の仕様を調べていて時間切れになりました。
あとで完成形にしたものの、提出したらWAでした。`(1<<i)`の直後の`% MOD`をつけていなかったためで、例によってリリースビルドのみオーバーフロー無視の罠にまたはまりました。
実戦中に5完できた可能性はかなり低かったと思います。

自分でライブラリを準備して臨むのもよいですが、利用できるクレートは利用することによって短くて可読性の良いコードを書く方がRustの精神には沿っていると思われ、petgraphも積極的に利用してくのが良いと考えます。