LoginSignup
4

posted at

Rust petgraph で競技プログラミングのグラフ理論系頻出アルゴリズムを学ぶ

はじめに

AtCoder に参加して 3か月ほど経ちました。グラフ問題を見ると身構えてしまいます。

Rust 初心者の AtCoder でよく使う言語機能とライブラリー覚え書き - Qiita で Union Find のために petgraph を使いました。せっかくだからとその流れで、petgraph のグラフ系のアルゴリズム関連の使い方を調べました。

想定読者

  • 競技プログラミングに参加している・興味がある
  • Rust の経験がある、または Rust 以外の何かのプログラミング言語を使える
  • グラフ理論系のアルゴリズムを眺めたい、使いたい
  • DP (動的計画法)、スタック、キュー、優先度付きキューを使ったことがある

本記事で行うこと

  • 競技プログラミングで、グラフ理論系で頻出するアルゴリズムを短く紹介
    • 最短経路探索: ダイクストラ法、ベルマン-フォード法、ワーシャル-フロイド法
    • 最小全域木問題: Union Find、クラスカル法
    • その他: ループ判定、二部グラフ判定、トポロジカルソート
    • AtCoder で言うと緑あたりの難易度まで (書いている人が緑色です)
  • ライブラリーなしにその場でぱっと書く場合のコード
  • petgraph ライブラリーを使って短く書く場合のコード

本記事で行わないこと

petgraph とは

グラフアルゴリズムを提供する Rust のライブラリーです。AtCoder で version 0.5.0 を使用できます。

公式トップページに、ダイクストラ法の例が載っています。

use petgraph::graph::{NodeIndex, UnGraph};
use petgraph::algo::dijkstra;

#[test]
fn test() {
    // 無向グラフを作成する。 node は `i32` 型、edge は重みなどの関連情報を持たない。
    let g = UnGraph::<i32, ()>::from_edges(&[
        (1, 2), (2, 3), (3, 4),
        (1, 4)]);
    
    // `1` から `4` への最短経路を求める。エッジ間の重みはすべて 1 とする。
    let node_map = dijkstra(&g, 1.into(), Some(4.into()), |_| 1);
    assert_eq!(&1i32, node_map.get(&NodeIndex::new(4)).unwrap());
}

Node1 から Node4 への最短経路は (1, 4) だけを通るものです。 node_map.get(&NodeIndex::new(4)) で 1つ手前のノード番号が 1 と取れます。自力で実装するより、ライブラリーにお任せするのが簡単です。

本記事内で出てくるグラフの型はこの3つです。

型名 説明
Graph 汎用の隣接グラフ
DiGraph 有向グラフ (GraphDirected パラメーターを付けたもの)
UnGraph 無向グラフ (GraphUndirected パラメーターを付けたもの)

Qiita でほかの方が書かれている記事です。

AtCoder での petgraph の使い方

2022年7月時点の、AtCoder での petgraph の使い方です。

Cargo.toml
# グラフ
petgraph = "=0.5.0"

# 挿入順を保持するhash table
indexmap = "=1.3.2"

最短経路探索

有向グラフの最短経路を求める3つの方法です。

アルゴリズム 問題の種類 エッジの重み 計算量
ダイクストラ法 単一始点最短経路問題 非負値 $O((E+V)logV)$
ベルマン-フォード法 単一始点最短経路問題 実数 $O(EV)$
ワーシャル-フロイド法 全点対最短経路問題 実数 $O(V^3)$

ここでは負のループは現れないものとします。負のループがあると、負のループをぐるぐる回っていると無限に距離を小さくすることができてしまいます。 Node0 から Node3 への距離を測れません。

ダイクストラ法

問題と方針

Node0 から各ノードへの最短経路・最短距離を考えます。今回はすべてのエッジの重みを $1$ とします。

結果は次のようになります。点線は最短経路としては使いません。

優先度付きキューを使うと効率的に解けます:

  1. 優先度付きキュー heap を用意する。
  2. heap に始点 Node0 を距離 = 優先度 $0$ で入れる。
  3. heap の中で一番距離の短いものを取り出す。距離 $0$ の Node0 が最短経路確定。結果の配列に記録する。
  4. heap に、Node0 から辿れる Node1 を距離 = 優先度 $1$ で入れる。
  5. heap の中で一番距離の短いものを取り出す。距離 $1$ の Node1 が最短経路確定。結果の配列に記録する。
  6. heap に、Node1 から辿れる Node2 を距離 = 優先度 $2$ で、Node4 を距離 = 優先度 $2$ で入れる。

heap から距離順に取り出しているうちに、Node7 から Node4 のような遠回りの経路も現れますが、すでに Node4 は距離 $2$ で確定していますので無視されます。こうしてたどり終えると最短距離が求まります。

「優先度付きキューで最初にたどり着いた経路が最短経路」という前提で動いています。エッジの重みが負になるとこの前提が成り立たなくなります。その時は別の方法を使います。

そのまま書く

use std::{cmp::Reverse, collections::BinaryHeap};

const I64MAX: i64 = 0x7fff_ffff_ffff_ffff;

fn dijkstra(edges: &[(usize, usize)], node_count: usize, node_from: usize) -> Vec<i64> {
    let mut edges_from_nodes = vec![Vec::<usize>::new(); node_count];
    for (node_from, node_to) in edges.iter() {
        edges_from_nodes[*node_from].push(*node_to);
    }

    let weight = 1;
    let mut heap = BinaryHeap::<(Reverse<i64>, usize)>::new();
    heap.push((Reverse(0), node_from));

    let mut shortest = vec![I64MAX; node_count];
    while let Some(x) = heap.pop() {
        let (score, node) = ((x.0).0, x.1);
        if shortest[node] > score {
            shortest[node] = score;
            for node_to in &edges_from_nodes[node] {
                heap.push((Reverse(score + weight), *node_to));
            }
        }
    }
    shortest
}

#[test]
fn test() {
    let edges = [
        (0, 1),
        (1, 2),
        (1, 4),
        (2, 3),
        (2, 7),
        (3, 0),
        (4, 5),
        (5, 6),
        (6, 7),
        (7, 4),
        (8, 5),
    ];
    let node_count = 9;

    let shortest = dijkstra(&edges, node_count, 0);
    assert_eq!(shortest, [0, 1, 2, 3, 2, 3, 4, 3, I64MAX]);
}

petgraph を使う

use petgraph::algo::dijkstra;
use petgraph::prelude::*;
use std::collections::HashMap;

#[test]
fn test() {
    let mut graph: DiGraph<(), ()> = DiGraph::new();
    let mut nodes = Vec::new();
    for _ in 0..9 {
        nodes.push(graph.add_node(()));
    }
    graph.extend_with_edges(&[
        (0, 1),
        (1, 2),
        (1, 4),
        (2, 3),
        (2, 7),
        (3, 0),
        (4, 5),
        (5, 6),
        (6, 7),
        (7, 4),
        (8, 5),
    ]);

    let res = dijkstra(&graph, nodes[0], None, |_| 1);
    let expected_res: HashMap<NodeIndex, i64> = [
        (nodes[0], 0),
        (nodes[1], 1),
        (nodes[2], 2),
        (nodes[3], 3),
        (nodes[4], 2),
        (nodes[5], 3),
        (nodes[6], 4),
        (nodes[7], 3),
    ]
    .iter()
    .cloned()
    .collect();
    assert_eq!(res, expected_res);
}

グラフを作って dijkstra() 関数を呼ぶだけ。簡単です。
エッジに重みがある場合は、dijkstra(&graph, nodes[0], None, |e| e.weight); のようにラムダ式を与えると良いです。

ベルマン–フォード法

問題と方針

同じく Node0 から各ノードへの最短経路を考えます。

ノードがすべてで6個あります。Node0 から最短経路があるなら、最大でも1本少ない5つのエッジをたどればたどり着くはずです。エッジ数を1つずつ増やしながら最短距離を更新します。

エッジ数 Node0 Node1 Node2 Node3 Node4 Node5
0 0
1 0 2 4
2 0 2 3 4 5
3 0 2 3 4 5 6
4 0 2 3 4 5 6
5 0 2 3 4 5 6

こう見るとと普通の DP です。
DP で多次元配列を作るのをサボって、1つの配列をその場で書き換えても良いです。1回の操作で2本以上のエッジを進むこともありますが、先に進みすぎても結果は変わりません。

そのまま書く

const I64MAX: i64 = 0x7fff_ffff_ffff_ffff;

fn bellman_ford(edges: &[(usize, usize, i64)], node_count: usize, node_from: usize) -> Vec<i64> {
    let edge_count = edges.len();
    let mut scores = vec![I64MAX; node_count];
    scores[node_from] = 0;

    for _ in 1..edge_count {
        for (node_from, node_to, weight) in edges {
            if scores[*node_from] < I64MAX {
                scores[*node_to] = scores[*node_to].min(scores[*node_from] + *weight);
            }
        }
    }
    scores
}

#[test]
fn test() {
    let edges = vec![
        (0, 1, 2),
        (0, 3, 4),
        (1, 2, 1),
        (1, 5, 7),
        (2, 4, 5),
        (4, 5, 1),
        (3, 4, 1),
    ];
    let node_count = 6;

    let shortest = bellman_ford(&edges, node_count, 0);
    assert_eq!(shortest, [0, 2, 3, 4, 5, 6]);
}

petgraph を使う

petgraph の bellman_ford() はエッジの重みを実数で計算するようです。 f64 にしています。
最短経路を取るときの一つ手前のノードも簡単に分かります。

use petgraph::algo::bellman_ford;
use petgraph::graph::DiGraph;

#[test]
fn test() {
    let mut graph: DiGraph<(), f64> = DiGraph::new();
    let mut nodes = Vec::new();
    for _ in 0..6 {
        nodes.push(graph.add_node(()));
    }
    graph.extend_with_edges(&[
        (0, 1, 2.0),
        (0, 3, 4.0),
        (1, 2, 1.0),
        (1, 5, 7.0),
        (2, 4, 5.0),
        (4, 5, 1.0),
        (3, 4, 1.0),
    ]);

    let path = bellman_ford(&graph, nodes[0]);
    assert_eq!(
        path,
        Ok((
            vec![0.0, 2.0, 3.0, 4.0, 5.0, 6.0],
            vec![
                None,
                Some(nodes[0]), // node0 -> node1
                Some(nodes[1]), // node1 -> node2
                Some(nodes[0]), // node0 -> node3
                Some(nodes[3]), // node3 -> node4
                Some(nodes[4])  // node4 -> node5
            ]
        ))
    );
}

ワーシャル–フロイド法

問題と方針

ベルマン–フォード法と同じグラフです。今度は各ノードから各ノードへの最短経路を考えます。
ワーシャル-フロイド法は、単純な方法で、すべての始終点の組み合わせをまとめて解けます。
(ベルマン–フォード法で始点を Node0..Node5 それぞれについて回しても解けます。どちらを使うかはお好みで。)

解き方です:

  1. Node1 から Node2 などに直接移動する場合のエッジ重みを並べた行列を作ります。自分自身への移動は $0$、移動できないものは $∞$ とします。
  2. Node0 を途中経路にできる」ように表を更新します。例えば Node1 から Node0 経由で Node2 にたどる距離がもとの距離より短くなれば、表の値を書き換えます。
  3. Node0Node1 を途中経路にできる」ように表を更新します。同様に行います。
  4. この操作を繰り返し、「どのノードも途中経路にできる」ように表を更新します。これが答えとなります。

最初の状態:

from\to Node0 Node1 Node2 Node3 Node4 Node5
Node0 0 2 4
Node1 0 1 7
Node2 0 5
Node3 0 1
Node4 0 1
Node5 0

最後の状態:

from\to Node0 Node1 Node2 Node3 Node4 Node5
Node0 0 2 3 4 5 6
Node1 0 1 6 7
Node2 0 5 6
Node3 0 1 2
Node4 0 1
Node5 0

こちらも普通の DP です。

そのまま書く

k, i, j の3重ループが肝です、と言うかそれだけです。i, j, k の順に回すと違う答えになります。

const I64MAX: i64 = 0x7fff_ffff_ffff_ffff;

fn floyd_warshall(edges: &[(usize, usize, i64)], node_count: usize) -> Vec<Vec<i64>> {
    let mut m = vec![vec![I64MAX; node_count]; node_count];
    for i in 0..node_count {
        m[i][i] = 0;
    }
    for (node_from, node_to, weight) in edges {
        m[*node_from][*node_to] = *weight;
    }

    for k in 0..node_count {
        for i in 0..node_count {
            for j in 0..node_count {
                if m[i][j] < I64MAX && m[j][k] < I64MAX {
                    m[i][k] = m[i][k].min(m[i][j] + m[j][k]);
                }
            }
        }
    }
    m
}

#[test]
fn test() {
    let edges = [
        (0, 1, 2),
        (0, 3, 4),
        (1, 2, 1),
        (1, 5, 7),
        (2, 4, 5),
        (4, 5, 1),
        (3, 4, 1),
    ];
    let node_count = 6;

    let shortest = floyd_warshall(&edges, node_count);
    assert_eq!(
        shortest,
        [
            [0, 2, 3, 4, 5, 6],
            [I64MAX, 0, 1, I64MAX, 6, 7],
            [I64MAX, I64MAX, 0, I64MAX, 5, 6],
            [I64MAX, I64MAX, I64MAX, 0, 1, 2],
            [I64MAX, I64MAX, I64MAX, I64MAX, 0, 1],
            [I64MAX, I64MAX, I64MAX, I64MAX, I64MAX, 0],
        ]
    );
}

petgraph を使う

petgraph 0.6 から floyd_warshall() アルゴリズムが使えるようです。 2
petgraph 0.5 では残念ながら使えません。2022年7月時点の AtCoder で使えませんので省略します。

最小全域木問題

無向グラフが与えられます。すべてのノードを辿れるようにしたままエッジを減らしていき、残ったエッジの重みの合計がもっとも小さくなるグラフを作る、という問題です。

この場合はどの1本のエッジを除いてもループがなくなります。 $3$ の重みのエッジを消すと、合計が一番小さくなります。これが最小全域木です。ループがなくなり、エッジの数がノードの数 -1 になります。

クラスカル法、プリム法をよく使うようです。本記事ではクラスカル法と、その道具として Union Find を紹介します。

Union Find

グループが分かれているかの判定方法です。AtCoder 公式解説が詳しいです。

問題と方針

(Node0, Node2) が同じグループ、(Node0, Node6) が同じグループ、のように複数の組が与えられたときに、(Node0, Node3) が同じグループ化を判定します。

矢印を辿った一番先のノードが同じかどうかで判定します。ノードが浅いほど辿るのが速くなりますので、つなぎ方に工夫があります。詳しくは AtCoder 解説をどうぞ。

そのまま書く

AtCoder 解説の rank なし版をそのまま書きました。

fn find(uf: &[usize], x: usize) -> usize {
    if uf[x] == x {
        x
    } else {
        let x = uf[x];
        find(uf, x)
    }
}

fn find_mut(uf: &mut [usize], x: usize) -> usize {
    uf[x] = find(uf, uf[x]);
    uf[x]
}

fn union(uf: &mut [usize], x: usize, y: usize) -> bool {
    if x == y {
        return false;
    }
    let x = find_mut(uf, x);
    let y = find_mut(uf, y);
    if x == y {
        return false;
    }
    uf[y] = x;
    true
}

#[test]
fn test() {
    let count = 8;
    let mut uf = vec![0; count];
    for i in 0..count {
        uf[i] = i;
    }

    union(&mut uf, 0, 2); // 集合 [0, 2]
    union(&mut uf, 2, 4); // 集合 [0, 2, 4]
    union(&mut uf, 5, 6); // 集合 [5, 6]
    union(&mut uf, 0, 6); // 集合 [0, 2, 4, 5, 6]

    union(&mut uf, 1, 3); // 集合 [1, 3]

    assert_eq!(find(&uf, 4), find(&uf, 6)); // [4, 6] は同じ集合に属する
    assert_ne!(find(&uf, 3), find(&uf, 4)); // [3, 4] は異なる集合に属する
}

petgraph を使う

Rust 初心者の AtCoder でよく使う言語機能とライブラリー覚え書き - Qiita の通りです。

use petgraph::unionfind::UnionFind;

#[test]
fn test() {
    // UnionFind 
    let mut uf = UnionFind::<usize>::new(8);
    uf.union(0, 2); // 集合 [0, 2]
    uf.union(2, 4); // 集合 [0, 2, 4]
    uf.union(5, 6); // 集合 [5, 6]
    uf.union(0, 6); // 集合 [0, 2, 4, 5, 6]

    uf.union(1, 3); // 集合 [1, 3]

    assert_eq!(uf.find(4), uf.find(6)); // [4, 6] は同じ集合に属する
    assert_ne!(uf.find(3), uf.find(4)); // [3, 4] は異なる集合に属する
}

クラスカル法

問題と方針

ひとつながりのグラフになるように、重みの小さなエッジから順につないでいきます。

2つのードが同じグループになっていなことを Union Find で確認してからつなぎます。
同じグループの場合はつなぎません。ループになりますので。

重み エッジ 使う? グループ
5 Node0 - Node3 true [0, 3]
5 Node2 - Node4 true [0, 3], [2, 4]
6 Node3 - Node5 true [0, 3, 5], [2, 4]
7 Node0 - Node1 true [0, 1, 3, 5], [2, 4]
7 Node1 - Node4 true [0, 1, 2, 3, 4, 5]
8 Node1 - Node2 false [0, 1, 2, 3, 4, 5]
8 Node4 - Node5 false [0, 1, 2, 3, 4, 5]
9 Node1 - Node3 false [0, 1, 2, 3, 4, 5]
9 Node4 - Node6 true [0, 1, 2, 3, 4, 5, 6]
11 Node5 - Node6 false [0, 1, 2, 3, 4, 5, 6]
15 Node3 - Node4 false [0, 1, 2, 3, 4, 5, 6]

これでループがなくなりました。

そのまま書く

use std::{cmp::Reverse, collections::BinaryHeap};

// find(), find_mut(), union() は UnionFind と同じ。省略。

fn kruskal(edges: &[(usize, usize, i64)], node_count: usize) -> Vec<(usize, usize)> {
    let mut uf = vec![0; node_count];
    for i in 0..node_count {
        uf[i] = i;
    }

    let mut heap = BinaryHeap::<(Reverse<i64>, (usize, usize))>::new();
    for (node_from, node_to, weight) in edges {
        heap.push((Reverse(*weight), (*node_from, *node_to)));
    }

    let mut spanning_tree_edges = Vec::<(usize, usize)>::new();
    for (_, (node_from, node_to)) in heap {
        if find(&uf, node_from) != find(&uf, node_to) {
            union(&mut uf, node_from, node_to);
            spanning_tree_edges.push((node_from, node_to));
        }
    }
    spanning_tree_edges
}

#[test]
fn test() {
    let edges = [
        (0, 1, 7),
        (0, 3, 5),
        (1, 2, 8),
        (1, 3, 9),
        (1, 4, 7),
        (2, 4, 5),
        (3, 4, 15),
        (3, 5, 6),
        (4, 5, 8),
        (4, 6, 9),
        (4, 6, 11),
    ];
    let node_count = 7;

    let mut mst_edges = kruskal(&edges, node_count);
    mst_edges.sort();
    assert_eq!(mst_edges, [(0, 1), (0, 3), (1, 4), (2, 4), (3, 5), (4, 6)]);
}

petgraph を使う

use petgraph::algo::min_spanning_tree;
use petgraph::data::FromElements;
use petgraph::graph::UnGraph;

#[test]
fn test() {
    let graph = UnGraph::<(), i64>::from_edges(&[
        (0, 1, 7),
        (0, 3, 5),
        (1, 2, 8),
        (1, 3, 9),
        (1, 4, 7),
        (2, 4, 5),
        (3, 4, 15),
        (3, 5, 6),
        (4, 5, 8),
        (4, 6, 9),
        (5, 6, 11),
    ]);

    let mst = UnGraph::<_, _>::from_elements(min_spanning_tree(&graph));
    let mut edges: Vec<_> = mst
        .raw_edges()
        .iter()
        .map(|x| (x.source().index(), x.target().index()))
        .collect();
    edges.sort();
    assert_eq!(edges, [(0, 1), (0, 3), (1, 4), (2, 4), (3, 5), (4, 6)]);
}

その他アルゴリズム

いろいろあります。ここからは駆け足でコードを1つずつ載せます。

ループ判定

有向グラフ

use petgraph::{algo::is_cyclic_directed, graph::DiGraph};

#[test]
fn test() {
    let g: DiGraph<(), ()> = DiGraph::from_edges(&[(0, 1), (1, 2), (2, 0)]);
    assert_eq!(is_cyclic_directed(&g), true);

    let g: DiGraph<(), ()> = DiGraph::from_edges(&[(3, 4), (4, 5), (3, 5)]);
    assert_eq!(is_cyclic_directed(&g), false);
}

無向グラフ

use petgraph::{algo::is_cyclic_undirected, graph::UnGraph};

#[test]
fn test() {
    let g: UnGraph<(), ()> = UnGraph::from_edges(&[(0, 1), (1, 2), (2, 0)]);
    assert_eq!(is_cyclic_undirected(&g), true);

    let g: UnGraph<(), ()> = UnGraph::from_edges(&[(3, 4), (4, 5)]);
    assert_eq!(is_cyclic_undirected(&g), false);
}

二部グラフ判定

隣り合うノードがすべて異なるように 2色に塗り分けられるかどうかを判定します。

petgraph 0.5 では残念ながら is_bipartite_undirected() を使えません。
頂点数を倍にして、2つのグループに分かれているかという方法で調べます。

use petgraph::unionfind::UnionFind;

fn is_bipartite(edges: &[(usize, usize)], node_count: usize) -> bool {
    let mut uf = UnionFind::<usize>::new(node_count * 2);
    for &(a, b) in edges.iter() {
        uf.union(a, b + node_count);
        uf.union(a + node_count, b);
    }

    (0..node_count)
        .into_iter()
        .all(|i| uf.find(i) != uf.find(i + node_count))
}

#[test]
fn test() {
    let mut edges = vec![
        (0, 1),
        (0, 3),
        (0, 5),
        (2, 1),
        (2, 3),
        (4, 1),
        (4, 3),
        (4, 5),
    ];
    let node_count = 6;
    assert_eq!(is_bipartite(&edges, node_count), true);

    edges.push((0, 2));
    assert_eq!(is_bipartite(&edges, node_count), false);
}

トポロジカルソート

有向グラフの順番を満たすように直列に並べます。複数の解があり得ます。

use petgraph::{
    algo::{toposort, DfsSpace},
    graph::DiGraph,
};

#[test]
fn test() {
    let g: DiGraph<(), ()> = DiGraph::from_edges(&[
        (1, 4),
        (1, 6),
        (2, 7),
        (3, 4),
        (4, 5),
        (3, 7),
        (7, 0),
        (7, 5),
    ]);
    let mut space = DfsSpace::new(&g);
    let result = toposort(&g, Some(&mut space));
    assert!(result.is_ok());

    let nodes: Vec<_> = result.unwrap().iter().map(|x| (x.index())).collect();
    let expected_nodes = [
        [3, 2, 1, 7, 4, 0, 5, 6],
        [1, 2, 3, 4, 7, 0, 5, 6],
        [1, 3, 4, 2, 7, 6, 0, 5],
        [2, 3, 1, 4, 7, 6, 5, 0],
        [3, 2, 7, 1, 6, 4, 5, 0],
        [3, 2, 7, 0, 1, 4, 5, 6],
    ];
    assert!(expected_nodes.iter().any(|expected| nodes.eq(&expected)));
}

最後に

アルゴリズムを呼び出せば終了、というものを調べていたつもりでした。だんだんグラフ理論が分かってきて、辿り方も気になってきました。

最近このような調べものを楽しんでいます。代わりに過去問挑戦が全然進んでいません。

  1. indexmap のバージョン指定が必要らしいです。indexmap 最新を取得しようとして、 Rust 2021 はサポートしないので使えないというエラーが出ました。

  2. 日本語では「ワーシャル-フロイド」の順番で呼ぶことが多いです。英語では floyd_warshall の順が多いようです

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
What you can do with signing up
4