irisuinwl’s diary

サークル不思議(略)入巣次元の、数学や技術的なことを書きます。

ABC 139 Eが解けた話

戦いたい選手をキューで詰めていってマッチングしたら戦わせる。 対戦(i,j)をグラフとしてDFSという方法がeditorialに載っているが、pythonでそれを表現したらTLEした。 (pythonでグラフ構造を効率よく表現する方法を知らなさすぎる) この方法で普通に解けたと思ったけどwhileループ中で抜けた場合の初期化処理をしていなくてWAした。

from unittest.mock import patch

testcases = """3
2 3
3 1
1 2
""".split('\n')

def solve():
    N = int(input())
    tournament = [list(map(lambda x:int(x)-1, input().split())) for _ in range(N)]
    map_que_tournament = {i:[] for i in range(N)} # player iがjと戦いたがってる
    competition_num = -1
    que_player = list(range(N)) #それぞれ戦えるプレイヤーをqueueにつめていく
    
    while que_player:
        day_process_flag = False
        new_queue = [] # 次のステップ更新用のqueue
        # print(que_player)
        for player in que_player:
            if tournament[player]:
                # トーナメントの対戦希望表から抜き出す
                match_player = tournament[player].pop(0)
                # ぶじマッチング相手が見つかったら
                if map_que_tournament[match_player] and player == map_que_tournament[match_player][0]:
                    # print(player, match_player)
                    day_process_flag = True
                    new_queue.append(match_player)
                    new_queue.append(player)
                    map_que_tournament[match_player].pop()
                else:
                    map_que_tournament[player].append(match_player)

        # 戦える人がいた場合に一日進める
        if day_process_flag:
            if competition_num < 0:
                competition_num = 1
            else: competition_num += 1

        que_player = new_queue
        
    # 途中でやめた場合
    for t in tournament:
        if t:
            competition_num = -1
            break

    print(competition_num)

with patch('builtins.input', side_effect=testcases):
    solve()

TLEしたグラフDFSしたpythonコード Nを最大にすると8秒くらい最初の読み込みでかかってしまう

def solve():
    fix_index = lambda x,y: (x,y) if x < y else (y,x)
    N = int(input())
    tournament_graph = {} # N^2
    tournament_graph_dom = {} # N^2
    search_vertex = {}
    max_path = -1
    
    for i in range(N):
        vertex_pre = None
        for j in map(int, input().split()):
            j = j-1
            vertex = fix_index(i,j)
            if vertex not in tournament_graph: # vertex not in tournament_graph
                tournament_graph[vertex] = []
                tournament_graph_dom[vertex] = []
                search_vertex[vertex] = True
            if vertex_pre:
                tournament_graph[vertex_pre].append(vertex)
                tournament_graph_dom[vertex].append(vertex_pre)
                
            vertex_pre = vertex
    
    vertices = []
    for vertex in tournament_graph_dom:
        # 点が存在して、入ってくる辺が無いもの
        if search_vertex[vertex] and (not tournament_graph_dom[vertex]):
            vertices.append(vertex)
            search_vertex[vertex] = False
    
    while tournament_graph:
        # ここで探索する
        new_vertives = []
        if vertices:
            # print(vertices)
            if max_path < 0: max_path = 0
            for v in vertices:
                while tournament_graph[v]: 
                    cod = tournament_graph[v].pop(0)
                    tournament_graph_dom[cod].pop(0)
                    # つぎ検索する箇所を選定
                    if search_vertex[cod] and (not tournament_graph_dom[cod]):
                        new_vertives.append(cod)
                        search_vertex[cod] = False
                    
            max_path += 1
        else:
            # DAG ならばdom(v) = \emptysetとなるvが存在しない
            break
            
        vertices = new_vertives
    
    print(max_path)
    
solve()