【競プロノート】グラフを Rust で実装してみる

先日の Rust LT 会でグラフを実装してみる話がありました。そこではおそらく隣接リストを説明していたと思うのですが、そういえば Rust で実装してみたことはなかったなと思ったので、実装してみました。

『みんなのデータ構造』という本を参考にして実装しています。12章にグラフの章があります。何言語かで書かれたサンプルコードも存在します。

今回は、本書に掲載されていた下記のような「典型的な操作」を実装しました。今回はグラフ G = (V, E) を考えるものとします。一般的な話通りで、V は頂点で、E は辺です。

• addEdge(i, j): 辺 (i, j) を E に加える
• removeEdge(i, j): 辺 (i, j) を E から除く
• hasEdge(i, j): (i, j) ∈ E かどうかを調べる
• outEdges(i): (i, j) ∈ E を満たす整数整数 j のリストを返す
• inEdges(i): (j, i) ∈ E を満たす整数整数 j のリストを返す

隣接行列

いきなりですが、隣接行列の実装は意外と大変です。Rust は、ポインタを使った操作 (参照先を書き換えてゴニョゴニョする操作) をしようとすると少々めんどくさくなるように言語が設計されています。隣接行列は行列 (2次元の配列) をもったデータ構造なのですが、その行列を状態とみなして書き換えていく操作を発生させるため、Rust 的にはとても面倒な実装になります。

今回実装した隣接行列は、(i, j) ∈ E のときに true となり、そうでない場合は false となる要素が含まれている行列です。

速度面などの細かい話を気にせず、RefCell を使用して実装してみました。ちなみに、std::mem::replace でも実装できると思いますし、実際 RefCell#replace はそれを使用して実装されているため、もしかすると RefCell を使用するのは大げさだったかもしれません。

use std::cell::RefCell;

pub struct AdjacencyMatrix {
    dimension: usize,
    matrix: Vec<Vec<RefCell<bool>>>,
}

impl AdjacencyMatrix {
    pub fn new(dimension: usize) -> AdjacencyMatrix {
        let mut matrix = vec![];
        for _ in 0..dimension {
            let mut tmp = vec![];
            for _ in 0..dimension {
                tmp.push(RefCell::new(false));
            }
            matrix.push(tmp);
        }

        AdjacencyMatrix { dimension, matrix }
    }

    pub fn dimension(&self) -> usize {
        self.dimension
    }

    fn get(&self, i: usize, j: usize) -> &RefCell<bool> {
        self.matrix.get(i).unwrap().get(j).unwrap()
    }

    pub fn add_edge(&self, i: usize, j: usize) {
        self.get(i, j).replace(true);
    }

    pub fn remove_edge(&self, i: usize, j: usize) {
        self.get(i, j).replace(false);
    }

    pub fn has_edge(&self, i: usize, j: usize) -> bool {
        *self.get(i, j).borrow()
    }

    pub fn in_edges(&self, i: usize) -> Vec<usize> {
        let mut edges = vec![];
        for j in 0..self.dimension {
            if *self.get(j, i).borrow() {
                edges.push(j);
            }
        }
        edges
    }

    pub fn out_edges(&self, i: usize) -> Vec<usize> {
        let mut edges = vec![];
        for j in 0..self.dimension {
            if *self.get(i, j).borrow() {
                edges.push(j);
            }
        }
        edges
    }
}

隣接リスト

逆に、隣接リストはスムーズに実装できます。これは、単純にエッジを追加する際にリストの末尾に値を入れていけばよいためです。

pub struct AdjacencyList {
    dimension: usize,
    list: Vec<Vec<usize>>,
}

impl AdjacencyList {
    pub fn new(dimension: usize) -> AdjacencyList {
        AdjacencyList {
            dimension,
            list: vec![vec![]],
        }
    }

    pub fn dimension(&self) -> usize {
        self.dimension
    }

    pub fn add_edge(&mut self, i: usize, j: usize) {
        if i > self.dimension() {
            panic!("add: greater than dimension");
        }

        match self.list.get_mut(i) {
            Some(v) => v.push(j),
            None => self.list.insert(i, vec![j]),
        }
    }

    pub fn remove_edge(&mut self, i: usize, j: usize) {
        match self.list.get_mut(i) {
            Some(v) => v.retain(|v0| *v0 != j),
            None => (),
        }
    }

    pub fn has_edge(&self, i: usize, j: usize) -> bool {
        match self.list.get(i) {
            Some(v) => v.contains(&j),
            None => false,
        }
    }

    pub fn in_edges(&self, i: usize) -> Vec<usize> {
        let mut edges = vec![];
        for j in 0..self.dimension() {
            if let Some(list) = self.list.get(j) {
                if list.contains(&i) {
                    edges.push(j);
                }
            }
        }
        edges
    }

    pub fn out_edges(&self, i: usize) -> Vec<usize> {
        match self.list.get(i) {
            Some(v) => v.to_vec(),
            None => vec![],
        }
    }
}

ポイントは、隣接リストの場合は初期状態の Vec<Vec<usize>> で、たとえば行を取得したい場合に、行が None として返ってきてしまうパターンが考えられるため、その点について留意しなが実装する必要があったという点です。Vector 関連でもっといい操作があれば知りたいところですが、思いつく限りでは None がどうしても発生してしまうため、都度パターンマッチして慎重に取り出す操作を行っています。

もっとも単純なグラフの実装がこれでわかりました。