E - Souvenir

[hxrxchang]

https://atcoder.jp/contests/abc286/tasks/abc286_e

問題概要

N個の都市があるのと、各都市から各都市への直行便があるかどうかと、各都市のお土産の価値の大きさが与えられる。 以下のクエリがQ個与えられるので答えよ。

• 都市UからVまで行くにあたり、使う直行便の数が最小、かつ、お土産の価値が最大になるような経路の直行便の数とお土産の価値の総和を求めよ。

解き方

任意の2頂点間の最短距離を求めたいのでワーシャルフロイド法を使う。おさらいだがダイクストラ法は単一始点の各頂点への最短距離なので今回は使えない。 移動コストは1でOKで、同時に獲得したお土産の価値の総和を記録しておく。 移動コストがより小さいルートが分かり最短距離を更新するときに、同時にそのルートでの獲得お土産価値も更新する。 移動コストが同じルートが出たら、お土産の価値が大きかったらそれを更新する。

func solve() {
    n := getInt()
    a := getInts()

    graph := make([][]int, n)
    for i := 0; i < n; i++ {
        in := strToSlice(getStr(), "")
        graph[i] = make([]int, n)
        for j, v := range in {
            if v == "Y" {
                graph[i][j] = a[j]
            }
        }
    }

    graph2 := warshallFloyd(graph)

    q := getInt()
    for i := 0; i < q; i++ {
        in := getInts()
        u, v := in[0]-1, in[1]-1
        if graph2[u][v].steps == BIGGEST {
            fmt.Println("Impossible")
        } else {
            fmt.Println(graph2[u][v].steps, a[u] + graph2[u][v].value)
        }
    }
}

type Node struct {
    steps int
    value int
}

func warshallFloyd(graph [][]int) [][]Node {
    n := len(graph)
    dist := make([][]Node, n)
    for i := range dist {
        dist[i] = make([]Node, n)
        for j := range dist[i] {
            if i == j {
                dist[i][j] = Node{0, 0}
            } else if graph[i][j] == 0 {
                dist[i][j] = Node{BIGGEST, 0}
            } else {
                dist[i][j].steps = 1
                dist[i][j].value = graph[i][j]
            }
        }
    }

    for i := 0; i < n; i++ {
        for j := 0; j < n; j++ {
            for k := 0; k < n; k++ {
                if dist[j][k].steps > dist[j][i].steps + dist[i][k].steps {
                    dist[j][k].steps = dist[j][i].steps + dist[i][k].steps
                    dist[j][k].value = dist[j][i].value + dist[i][k].value
                } else if dist[j][k].steps == dist[j][i].steps + dist[i][k].steps && dist[j][k].value < dist[j][i].value + dist[i][k].value {
                    dist[j][k].value = dist[j][i].value + dist[i][k].value
                }
            }
        }
    }

    return dist
}

https://atcoder.jp/contests/abc286/submissions/55131510