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

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

問題は、AtCoder ABC201E です。
アルゴリズムそのものについては公式サイトを見てください。

---

useなど

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};

---

入力を読み込んでグラフを構成

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で扱うなら、

-	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

	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[_]の記法を用いない場合、

-			g[ny] = g[nx] ^ g[edge];
+			let mut nyw = g.node_weight_mut(ny);
+			*nyw = g.node_weight(nx) ^ g.edge_weight(edge);

などとなります。

---

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

	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]で頂点の重みにアクセスします。

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

	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も積極的に利用してくのが良いと考えます。