ワーシャル-フロイド法について勉強する - どこにでもいるSEの備忘録

AtCoderやってたら、ワーシャル–フロイド法を使う問題が出てきて、知らないアルゴリズムだったので勉強してみた次第です。

ワーシャル–フロイド法 is なに？
例
実際にやってみる
参考文献
感想

ワーシャル–フロイド法 is なに？

wikipediaによるとこんな感じに書いてありました。

ワーシャル–フロイド法（英: Floyd–Warshall Algorithm）は、重み付き有向グラフの全ペアの最短経路問題を多項式時間で解くアルゴリズムである。名称は考案者であるスティーブン・ワーシャル（英語版）とロバート・フロイドにちなむ（2人はそれぞれ独立に考案）。フロイドのアルゴリズム、ワーシャルのアルゴリズム、フロイド–ワーシャル法とも呼ばれる。ワーシャル–フロイド法 - Wikipedia

例

重み付き有向グラフに関して、全ペアの最短経路問題を解くことを考えます。例えば、下記のような状況があったとします。

f:id:nogawanogawa:20210705004852j:plain

ここで、結合された場合のコスト>0とし、結合されていない場合をinf、始点と終点が同じノードであるケースはコスト0としています。

その際に、グラフ内のノードのペア全てに関して最短経路の距離を求めることとします。これを解くことを考えると、経由するノードに着目して、下記のような３重ループによって求めることができます。

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

func warshall_floyd(A [][]int, n int) { // nは頂点数
    for i := 0; i < n; i++ { // 経由する頂点
        for j := 0; j < n; j++ { // 開始頂点
            for k := 0; k < n; k++ { // 終了頂点
                A[i][j] = min(A[i][j], A[i][k]+A[k][j])
            }
        }
    }
}

これにより、始点 - 終点のペアに対する最短経路が二次元配列に書き込まれることになります。

このようなアルゴリズムがワーシャル–フロイド法のようです。

実際にやってみる

実際に、できなかった問題についてやってみます。

atcoder.jp

基本的な考え方としては上のワーシャル–フロイド法となっています。異なる点としては、経由する頂点による場合分けが発生している点です。それを踏まえて考えると下記のようなコードでうまくいくはずです。

クリックして展開

package main

import (
    "bufio"
    "fmt"
    "os"
    "strconv"
)

var sc = bufio.NewScanner(os.Stdin)

func nextInt() int {
    sc.Scan()
    i, e := strconv.ParseInt(sc.Text(), 10, 64)
    if e != nil {
        panic(e)
    }
    return int(i)
}

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

func warshall_floyd(A [][]int, B [][]int, n int, k int) int { // nは頂点数
    var ans int = 0

    for i := 0; i < n; i++ { // 開始頂点
        for j := 0; j < n; j++ { // 終了頂点
            A[i][j] = min(B[i][j], B[i][k]+B[k][j])
            if A[i][j] < inf {
                ans += B[i][j]
            }
        }
    }

    return ans
}

const inf int = 1e18

func main() {

    sc.Split(bufio.ScanWords)

    var N, M int
    fmt.Scanf("%d %d", &N, &M)

    A := make([][]int, N)
    B := make([][]int, N)

    for i := 0; i < N; i++ {
        A[i] = make([]int, N)
        B[i] = make([]int, N)
    }

    for i := 0; i < N; i++ {
        for j := 0; j < N; j++ {
            A[i][j] = inf
        }
        A[i][i] = 0
    }

    var start, goal, cost int
    for i := 0; i < M; i++ {
        start = nextInt()
        goal = nextInt()
        cost = nextInt()
        A[start-1][goal-1] = cost
    }

    var sum int = 0
    for i := 0; i < N; i++ {
        copy(B, A)
        sum += warshall_floyd(A, B, N, i)
    }

    fmt.Printf("%d\n", sum)
}