SHIROBAKO大好き人間のブログ

SHIROBAKOが好きなエンジニアによる技術ブログ

AtCoder Beginner Contest 061 問題D

Atcoderの問題を解いていたら面白いなあと思う問題があったので書いておきます。

問題

AtCoder Beginner Contest 061 問題D

 N 頂点  M 辺の重み付き有向グラフがあります。  i (1 \le i \le M) 番目の辺は 頂点  a_i から 頂点  b_i を重み  c_i で結びます。 このグラフと駒を利用して、次の1人ゲームを行います。

最初、駒を頂点  1 に置いて、プレイヤーのスコアを  0 とします。 プレイヤーは、次の条件で駒を繰り返し移動させることができます。

頂点  a_i に駒があるとき、 i 番目の辺を利用して頂点  b_i に移動する。移動後にプレイヤーのスコアが  c_i 加算される。 頂点  N に駒があるときのみ、ゲームを終了できます。 なお、与えられる有向グラフの上で頂点  1 から頂点  N に移動できることは保障されています。

プレイヤーがゲーム終了時のスコアを出来るだけ大きくするような行動を取ったとき、ゲーム終了時のスコアはいくつになるでしょうか? ゲーム終了時のスコアをいくらでも大きくできる場合は inf と出力してください。

(自分の)解答に至るまでのプロセス

ゲームと聞くとややこしく感じますが、要するに頂点  1から頂点 Nまでの重みが最大となる経路(最長経路)を求めればよさそうです。

infのパターンはややこしそうなので、とりあえず今は考えないでおきましょう。

ただ、ダイクストラなどの経路に関するアルゴリズムは最短経路を求めるアルゴリズムなので、これらをそのまま適用するのはダメそうです。

そこで、与えられたグラフの重みの符号を反転させます。

「もとのグラフにおける最長経路=符号を反転させたグラフにおける最短経路」なので、これで答えが求まりそうです。

あとの説明のためにもとのグラフを G、符号を反転させたグラフを G_{rev}と表記します。

符号を反転させたグラフ G_{rev}は負の重みを含みます。

ダイクストラは負の重みがある場合は使えないので、今回はベルマンフォード法を用います。

負の重みと聞くと難しそうに感じますが、アルゴリズムの実装はダイクストラと同じくらいの難易度です。

また、ベルマンフォード法は最短経路を求めるだけでなく、↓のノード 1とノード 2のように負の重みの閉路があるかどうかも検出することができます。

f:id:phoro3:20170516222354p:plain

ところで、このグラフの符号をもとに戻すとあることに気づきます。

f:id:phoro3:20170516223019p:plain

これ、ノード 1とノード 2をひたすら往復してから 3にゴールしたらスコアが無限大になりますね。

つまり、問題文中のスコアがinfになるのはベルマンフォード法において負の重みの閉路が検出されたパターンだったわけです。

この考えに基づいて書いたコードが以下のコードです。

#coding:utf-8
 
import sys
from collections import defaultdict
 
if __name__ == "__main__":
    N, M = map(int, input().split(" "))
 
    V = set()
    E = []
    for i in range(M):
        a, b, c = map(int, input().split(" "))
        V.add(a)
        V.add(b)
        E.append((a, b, -c))
 
    dist = defaultdict(lambda :10**12 + 1)
    dist[1] = 0
 
 
    for i in range(1, len(V)):
        for edge in E:
            u = edge[0]
            v = edge[1]
 
            if dist[v] > dist[u] + edge[2]:
                dist[v] = dist[u] + edge[2]
 
    is_cycle = False
    for edge in E:
        u = edge[0]
        v = edge[1]
        if dist[u] + edge[2] < dist[v]:
            is_cycle = True
 
    if is_cycle:
        print ("inf")
    else:
        print (-dist[N])

負の重みの閉路が検出されたらinf、それ以外はノード 1からノード Nの経路の長さを出力しています。

しかし、これを提出したらWAになってしまいました。

一部のテストケースのみでWAとなっていたので、考慮しきれていないパターンがあるようです。

何がダメだったのか?

結論から言うと、こういったパターンが抜けてました。(ノード 1がスタートでノード 4がゴールとする)

f:id:phoro3:20170516224006p:plain

つまり、閉路は存在するけれど、閉路から直接ゴールに向かうことができないパターンです。

この場合のスコアはinfではなく、ノード 1からノード 4に直接向かった場合の -2となります。

これを考慮したソースコードがこちら。

#coding:utf-8
 
import sys
from collections import defaultdict
 
if __name__ == "__main__":
    N, M = map(int, input().split(" "))
 
    V = set()
    E = []
    for i in range(M):
        a, b, c = map(int, input().split(" "))
        V.add(a)
        V.add(b)
        E.append((a, b, -c))
 
    dist = defaultdict(lambda :10**20 + 1)
    pre = {}
    dist[1] = 0
 
 
    for i in range(1, len(V)):
        for edge in E:
            u = edge[0]
            v = edge[1]
 
            if dist[v] > dist[u] + edge[2]:
                dist[v] = dist[u] + edge[2]
                pre[v] = u
 
    is_cycle = False
    cycle_node = []
    for edge in E:
        u = edge[0]
        v = edge[1]
        if dist[u] + edge[2] < dist[v]:
            is_cycle = True
            cycle_node.append(u)
 
    is_inf = False
    if is_cycle:
        route_node = N
        while route_node != 1:
            route_node = pre[route_node]
            if route_node in cycle_node:
                is_inf = True
                break
 
 
    if is_inf:
        print ("inf")
    else:
        print (-dist[N])

負の閉路が検出されたらその閉路を構成するノードをcycle_nodeに追加します。

そして、ベルマンフォード法の適用が終わった後に最短経路を確認し、最短経路中に閉路を構成するノードがあればinf、それ以外はスタートからゴールまでの経路の長さを出力しています。

こうすることで無事にACになりました!

解説の方にはもう少し賢いやり方も乗っているので興味がある方は見てみてください。