#1167006 (PyPy3) No.3561 Collect KCPC

• 問題: No.3561 Collect KCPC
• ユーザー: みうね
• 提出日時: 2026-05-28 11:26:34
• 言語: PyPy3 (7.3.17)
• 結果: AC
• 実行時間: 3,419 ms / 6,000 ms
• コード長: 4,088 bytes
• コンパイル時間: 686 ms
• コンパイル使用メモリ: 85,632 KB
• 実行使用メモリ: 208,244 KB
• 最終ジャッジ日時: 2026-05-29 18:48:39
• 合計ジャッジ時間: 34,679 ms
• ジャッジサーバーID （参考情報）: judge1_1 / judge2_1

サブタスク	配点	結果
部分点1	10 %	AC * 15
部分点2	20 %	AC * 15
部分点3	20 %	AC * 13
部分点4	50 %	AC * 51
合計	100 点

ソースコード

# Geminiで翻訳

import sys
import heapq

def solve():
    # 入力を一括で読み込み
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    N = int(input_data[0])
    M = int(input_data[1])
    
    # グラフ構築
    graph = [[] for _ in range(N)]
    idx = 2
    for _ in range(M):
        u = int(input_data[idx]) - 1
        v = int(input_data[idx+1]) - 1
        c = int(input_data[idx+2])
        graph[u].append((v, c))
        idx += 3
        
    S = input_data[idx]
    
    # 文字列比較のオーバーヘッドを無くすため、事前に整数へマッピング
    # K: 0, C: 1, P: 2, その他: -1
    node_type = [-1] * N
    for i in range(N):
        if S[i] == 'K':
            node_type[i] = 0
        elif S[i] == 'C':
            node_type[i] = 1
        elif S[i] == 'P':
            node_type[i] = 2

    INF = float('inf')
    
    # 状態配列を1次元化 (サイズ: N * 5)
    # index = u * 5 + s
    # リストのsort()を避けるため、1位(c1, id1)と2位(c2, id2)をマニュアルで管理
    size = N * 5
    dist_c1 = [INF] * size
    dist_id1 = [-1] * size
    dist_c2 = [INF] * size
    dist_id2 = [-1] * size

    # リスト操作・ソートのオーバーヘッドを排除した更新関数
    def update(idx, c, id_val):
        c1 = dist_c1[idx]
        id1 = dist_id1[idx]
        
        # すでに1位として登録されているIDの場合
        if id_val == id1:
            if c < c1:
                dist_c1[idx] = c
                return True
            return False
            
        c2 = dist_c2[idx]
        id2 = dist_id2[idx]
        
        # すでに2位として登録されているIDの場合
        if id_val == id2:
            if c < c2:
                # 1位を追い抜いた場合の入れ替え
                if c < c1:
                    dist_c1[idx], dist_c2[idx] = c, c1
                    dist_id1[idx], dist_id2[idx] = id_val, id1
                else:
                    dist_c2[idx] = c
                return True
            return False
            
        # 新しいIDの場合
        if c < c1:
            dist_c2[idx], dist_id2[idx] = c1, id1
            dist_c1[idx], dist_id1[idx] = c, id_val
            return True
        elif c < c2:
            dist_c2[idx], dist_id2[idx] = c, id_val
            return True
            
        return False

    # 優先度付きキュー (cost, u, state, id)
    pq = [(0, 0, 0, -1)]
    dist_c1[0] = 0
    dist_id1[0] = -1
    
    while pq:
        cost, u, s, id_val = heapq.heappop(pq)
        
        idx = u * 5 + s
        
        # 枝刈り (Lazy Deletion)
        if dist_id1[idx] == id_val:
            if dist_c1[idx] < cost: continue
        elif dist_id2[idx] == id_val:
            if dist_c2[idx] < cost: continue
        else:
            continue
            
        # 1. 隣接頂点への移動
        for v, edge_cost in graph[u]:
            nxt_idx = v * 5 + s
            nxt_cost = cost + edge_cost
            if update(nxt_idx, nxt_cost, id_val):
                heapq.heappush(pq, (nxt_cost, v, s, id_val))
                
        # 2. 状態の遷移 (K -> C -> P -> C)
        nt = node_type[u]
        if nt != -1:
            ns = -1
            nid = id_val
            
            if s == 0 and nt == 0:
                ns = 1
            elif s == 1 and nt == 1:
                ns = 2
                nid = u
            elif s == 2 and nt == 2:
                ns = 3
            elif s == 3 and nt == 1 and u != id_val:
                ns = 4
                
            if ns != -1:
                nxt_idx = u * 5 + ns
                if update(nxt_idx, cost, nid):
                    heapq.heappush(pq, (cost, u, ns, nid))
                    
    # 答えの計算 (状態4に到達している頂点の中での最小コスト)
    ans = INF
    for i in range(N):
        c = dist_c1[i * 5 + 4]
        if c < ans:
            ans = c
            
    print(-1 if ans == INF else ans)

if __name__ == '__main__':
    solve()