オットセイの経営日誌

データサイエンス系ベンチャーを経営してます。経営のこと、趣味のことつぶやきます。

DAG(有向非循環グラフ)に対する最長経路問題(AtCoder Beginner Contest 139より)

最近はKaggleを優先していてなかなか競プロの時間を取れないのですが、AtCoderのABCには参加しています。

で、1週間遅れになってしまいましたが、先日のABC139でDAGの問題が出ました。
DAGは今後も重要そうな気がして、YouTube公式解説のC++のコードを読み解き、Python化しました。

問題内容・解法

E - League

atcoder.jp

解法の考え方については以下の公式解説が非常に分かりやすいです。

www.youtube.com

簡単に紹介しておくと、問題で与えられたデータは、リーグ戦をイメージして、各選手について対戦したい選手の順番が決まっているというものですが、これを試合idに読み替えます。

f:id:mhiro216:20190908141016p:plain

すると、これはDAG(有向非循環グラフ)に対する最長経路問題であると読み替えることができます。

f:id:mhiro216:20190908141439p:plain

上の例では、頂点Aから頂点Eへの移動が最長経路となり、値は4、つまり試合日程を全てこなすのに必要な最小の日数は4日である、ということになります。

DAGということはループが存在しないわけですが、この問題でなぜそう読み替えられるかというと、
「選手1はまず選手2と対戦したくて、選手2はまず選手3と対戦したくて、選手3はまず選手1と対戦したい」などというデータがあった場合、これを満たす試合日程を組むことはできないわけですが、これはグラフ上ではループとして検出できます。
問題文では条件を満たすように試合日程を組めなければ'-1'を返せと言っているので、ループが検出されたら'-1'を返すように実装することになります。

さて、以下がPythonによる実装になります。

import sys
sys.setrecursionlimit(10**6)
readline = sys.stdin.readline

MAXN = 1005
MAXV = MAXN*(MAXN-1)//2
to = [[] for _ in range(MAXV)] # 頂点間の辺の情報
id = [[[] for _ in range(MAXN)] for _ in range(MAXN)] # 試合のid=DAGの頂点番号
def toId(i,j): # 選手idから試合idを返す関数
    if i > j: # i,jが逆になっても試合idは同じ
        i,j = j,i
    return id[i][j]

visited = [False]*MAXV
calculated = [False]*MAXV
dp = [1]*MAXV # Vからスタートしたときの最長経路。頂点の個数ベースで経路の長さを数えるので、初期値は1

def dfs(v):
    if visited[v]:
        if not calculated[v]:
            return -1 # 計算が終わっていない頂点を2度訪れるのはループがあるということ
        return dp[v]
    visited[v] = True
    for u in to[v]: # 全ての辺をなめる
        res = dfs(u)
        if res == -1: return -1 # ループがあれば-1を返す
        dp[v] = max(dp[v], res+1)
    calculated[v] = True
    return dp[v]

def main():
    n = int(input())
    a = [[int(i) for i in readline().split()] for _ in range(n)]
    for i in range(n):
        for j in range(n-1):
            a[i][j] -= 1 # 選手idを0始まりに変換
    V = 0
    for i in range(n):
        for j in range(n):
            if i < j:
                id[i][j] = V
                V += 1 # 0から順に各試合にidを割り振る
    for i in range(n):
        for j in range(n-1):
            a[i][j] = toId(i,a[i][j]) # 選手idを試合id(頂点番号)に置き換える
        for j in range(n-2): # 頂点間の依存関係はn-2個
            to[a[i][j+1]].append(a[i][j])
    ans = 0
    for i in range(V):
        res = dfs(i)
        if res == -1:
            print('-1') # ループがあれば-1を返す(問題文の指示)
            return
        ans = max(ans, res)
    print(ans)
    return

if __name__ == '__main__':
    main()

ここで(大)問題は、Pythonでsubmitすると一部のテストケースでTLEになってしまいます(泣)。
リスト内包表記を入れるとかしてチマチマ高速化しようか悩んでいたところ、Pypyでsubmitしてみるという手があることを知りました。

[https://juppy.hatenablog.com/entry/2019/06/14/Python%E7%AB%B6%E6%8A%80%E3%83%97%E3%83%AD%E3%82%B0%E3%83%A9%E3%83%9F%E3%83%B3%E3%82%B0%E9%AB%98%E9%80%9F%E5%8C%96tips%28Python%E3%81%A7Atcoder%E3%82%92%E3%82%84%E3%82%8B%E9%9A%9B%E3%81%AB%E5%80%8B:embed:cite]

Pypyで提出すると、無事ACとなりました。
低級言語への乗り換え圧力を強く感じましたが、なんとか乗り切ったことにします。