Pythonで解く【初中級者が解くべき過去問精選 100 問】(056 - 059 最短経路問題：ダイクストラ法)

1. 目的

初中級者が解くべき過去問精選 100 問をPythonで解きます。
すべて解き終わるころに水色になっていることが目標です。

本記事は「056 - 059 最短経路問題：ダイクストラ法」です。

2. 総括

ダイクストラ法とワーシャルフロイド法についてはYouTubeのAbdul Bariのチャンネルを見て理解しました。
とても分かりやすく、動画を見てアルゴリズムを理解してから問題を解いたのでスムーズに進めることができました。

3. 本編

056 - 059 最短経路問題：ダイクストラ法

056. GRL_1_A - 単一始点最短経路

回答

import heapq
#--------------------------------------[入力受取]--------------------------------------#
INF = float('inf')
V, E, r = map(int, input().split())
graph = [[] for _ in range(V)]
for _ in range(E):
    s, t, d = map(int, input().split())
    graph[s].append((t, d))
#--------------------------------------[初期値]--------------------------------------#
dp = [INF] * V
visited = [-1] * V
dp[r] = 0
h = [(0, r)]
#--------------------------------------[探索開始]--------------------------------------#
while h:
    d, s = heapq.heappop(h) # costが小さい順に取り出す
    if visited[s] == 1:
        continue
    visited[s] = 1
    targets = graph[s]
    for target in targets:
        t, d = target
        if dp[t] > dp[s] + d:
            dp[t] = dp[s] + d
            heapq.heappush(h, (dp[t], t))
#--------------------------------------[回答]--------------------------------------#
for answer in dp:
    if answer == INF:
        print('INF')
    else:
        print(answer)

ヒープで小さい順から取り出します。
アルゴリズムは3.6 Dijkstra Algorithm - Single Source Shortest Path - Greedy Methodがわかりやすいです。
この動画で解説されている通りコードに落としてやると回答となります。

---

057. JOI 2008 予選 6 - 船旅

回答

import heapq

def cal_cost(graph, n, start, end):
    #--------------------------------------[初期値]--------------------------------------#
    dp = [INF] * (n + 1)
    visited = [-1] * (n + 1)
    dp[start] = 0
    h = [(0, start)]
    #--------------------------------------[探索開始]--------------------------------------#
    while h:
        _, s = heapq.heappop(h) # sはstart, eはend
        if visited[s] == 1:
            continue
        visited[s] = 1
        targets = graph[s]

        for target in targets:
            cost, e = target
            if dp[e] > dp[s] + cost:
                dp[e] = dp[s] + cost
                heapq.heappush(h, (dp[e], e))

    return dp[end]


if __name__ == "__main__":
    INF = float('inf')
    n, k = map(int, input().split()) # nは島の数、kは以下の入力の行数
    graph = [[] for _ in range(n+1)] # 1インデックス

    answer_list = []
    for _ in range(k):
        info_list = list(map(int, input().split()))

        if info_list[0] == 1:
            island_1, island_2, cost = info_list[1], info_list[2], info_list[3]
            graph[island_1].append((cost, island_2))
            graph[island_2].append((cost, island_1))
        else:
            start, end = info_list[1], info_list[2]
            answer = cal_cost(graph, n, start, end)
            answer_list.append(answer)

    for answer in answer_list:
        if answer == INF:
            print(-1)
        else:
            print(answer)

問題56を少しだけ変えた問題ですが、基本的には行っていることは同じです。
とくにcal_cost()の中身はほぼ同じなので、入力の受け取りだけ、気を付けてやれば解けます。

---

058. JOI 2016 予選 5 - ゾンビ島

回答

from collections import deque
import heapq
#--------------------------------------[初期値]--------------------------------------#
INF = float('inf')
N, M, K, S = map(int, input().split()) # N個の町、M本の道路、K個がゾンビに支配、ゾンビからS本以内は危険な町
P, Q = map(int, input().split()) # 危険じゃない場合はP円、危険な場合はQ円
zombie_towns = [int(input()) for _ in range(K)]

graph = [[] for _ in range(N+1)]
for _ in range(M):
    townA, townB = map(int, input().split())
    graph[townA].append((INF, townB))
    graph[townB].append((INF, townA))

#--------------------------------------[BFSでゾンビ町からの距離を探索]--------------------------------------#
# zombie_townsからS本以内の危険な町を探すためのBFS
visited = [INF] * (N+1)
q = deque()
for zombie_town in zombie_towns:
    q.append(zombie_town)
    visited[zombie_town] = 0

while q:
    start_town = q.popleft()

    for target in graph[start_town]:
        _, end_town = target

        if visited[end_town] != INF:
            continue
        q.append(end_town)
        visited[end_town] = visited[start_town] + 1

# zombie_townsからS本以内の危険な町についてsetとして記録しておく
cost_Q = set()
for i in range(1, N+1):
    if visited[i] <= S:
        cost_Q.add(i)

#--------------------------------------[heapqで最短距離を探索]]--------------------------------------#
# heapをまわす
dp = [INF] * (N+1)
visited2 = [-1] * (N+1)
dp[1] = 0
h = [(0, 1)]
answer = 0
while h:
    cost, s = heapq.heappop(h) # sはstart, eはend
    if visited2[s] == 1:
        continue
    if s in zombie_towns: # ゾンビがいる町は行かない
        continue
    visited2[s] = 1
    targets = graph[s]

    for target in targets:
        _, e = target

        if e in cost_Q: #  zombie_townsからS本以内の危険な町はQ, それ以外はP
            cost = Q
        else:
            cost = P

        if dp[e] > dp[s] + cost:
            dp[e] = dp[s] + cost
            heapq.heappush(h, (dp[e], e))
        if e == N: # 目的地がでてきたらそこでbreak
            answer = dp[s] # 答えは目的地の一つ前で記録したコスト
            break
    if answer != 0:
        break

print(answer)

実装が重いです。
ただ、やることを分解してやると、いままで解いてきた問題の組み合わせなので、理解はしやすいです。
問題は大きく二つに分けることができ、「ゾンビの町からの距離を求める」という問題と「最小コストの道のりを求める」という問題です。

「ゾンビの町からの距離を求める」ためにはBFSでゾンビからの最短距離を求めればよいです。
「最小コストの道のりを求める」ためには、heapqを使ってダイクストラ法を使えばよいです。

---

059. JOI 2014 予選 5 - タクシー

回答(TLE (pypyだとMLE))

from collections import deque
import heapq

#------------------[BFSで各点からそれぞれの目的地までの距離を求める]----------------------#
def bfs(start, graph):

    visited = [-1] * N
    q = deque()
    q.append(start)
    visited[start] = 0

    while q:
        s = q.popleft()
        targets = graph[s]
        for e in targets:

            if visited[e] != -1:
                continue
            visited[e] = visited[s] + 1
            q.append(e)

    return visited


if __name__ == "__main__":
    #------------------[入力]----------------------#
    INF = float('inf')

    N, K = map(int, input().split()) # N個の町、K本の道路
    costs = [INF] * N
    counts = [0] * N
    for i in range(N):
        C, R = map(int, input().split())
        costs[i] = C
        counts[i] = R

    graph = [[] for _ in range(N)]
    for _ in range(K):
        A, B = map(int, input().split())
        A -= 1
        B -= 1
        graph[A].append(B)
        graph[B].append(A)

    #------------------[graphの再構築]]----------------------#
    graph2 = [[] for _ in range(N)]
    for start in range(N):
        end_list = bfs(start, graph)

        for end, count in enumerate(end_list):
            if counts[start] < count: 
                continue
            if start == end:
                continue
            graph2[start].append((costs[start], end)) #回数制限内にいける箇所のみgraphに追加

    #------------------[heapqで最短距離]]----------------------#
    dp = [INF] * N
    visited = [-1] * N
    dp[0] = 0
    h = [(0, 0)]
    while h:
        _, s = heapq.heappop(h)
        if visited[s] == 1:
            continue
        visited[s] = 1
        targets = graph2[s]

        for target in targets:
            cost, e = target
            if dp[e] > dp[s] + cost:
                dp[e] = dp[s] + cost
                heapq.heappush(h, (dp[e], e))

    print(dp[N-1])

上記コードだと、5ケースのうち2つがTLEになります。pypyで提出すると5ケースのうち1つだけMLEとなります。

大きな考え方は問題58と似ていて、 BFSを行ってからダイクストラ法を行います。

方針を書くと、
1. すべての町を始点として、graphをもとにそれぞれの町に何手で行けるかをbfsで算定
2. この手数のリストをend_listに入れておく
3. end_listと入力値Rを記録したcountsを使ってgraph2を作成します
4. graph2を作れれば、あとはダイクストラ法を行うだけ

です。

問題文であたえられるgraphはbfsに使い、
再構築したgraph2はダイクストラ法で使う、
という点が、この問題の面白いところでしょうか。

---