AT274の技術だったりなかったり

あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

CODE FESTIVAL 2014 Easy C - 身体バランス

問題原文

atcoder.jp

解法

 A→  x → B みたいなのがあるときには、真ん中を軸として考えると嬉しいことが多いんですが、今回は真ん中を全部調べる必要があります。

ということでダイクストラ S からの最短距離と  T からの最短距離を調べておけば、各  x について高速に判定ができます。

連結グラフではないのでちょっと注意。

実装

N, M = map(int, input().split())
S, T = map(int, input().split())
S, T = S - 1, T - 1
G = [[] for _ in range(N)]
for i in range(M):
    x, y, d = map(int, input().split())
    x, y = x - 1, y - 1
    G[x].append((y, d))
    G[y].append((x, d))


def dijkstra(graph, start, inf=float('inf')):
    import heapq
    n = len(graph)
    distances = [inf] * n
    distances[start] = 0
    visited = [False] * n

    # 距離・頂点
    hq = [(0, start)]
    while hq:
        dist, fr = heapq.heappop(hq)
        visited[fr] = True

        for to, cost in graph[fr]:
            new_dist = distances[fr] + cost
            if (visited[to]) or (distances[to] <= new_dist):
                continue

            distances[to] = new_dist
            heapq.heappush(hq, (new_dist, to))

    return distances


S_distances = dijkstra(G, S)
T_distances = dijkstra(G, T)
for u in range(N):
    if S_distances[u] == T_distances[u] and S_distances[u] != float('inf'):
        print(u + 1)
        break
else:
    print(-1)

感想

いつかPASTで出そうだなぁ(こなみ)