結果

問題 No.2431 Viral Hotel
ユーザー lam6er
提出日時 2025-04-09 21:02:54
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 305 ms / 2,000 ms
コード長 3,309 bytes
コンパイル時間 950 ms
コンパイル使用メモリ 81,844 KB
実行使用メモリ 146,396 KB
最終ジャッジ日時 2025-04-09 21:05:00
合計ジャッジ時間 12,106 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 4
other AC * 42
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
from collections import defaultdict

def main():
    input = sys.stdin.read().split()
    ptr = 0
    N = int(input[ptr]); ptr +=1
    K = int(input[ptr]); ptr +=1
    M = int(input[ptr]); ptr +=1
    P_val = int(input[ptr]); ptr +=1
    
    adj = defaultdict(list)
    for _ in range(M):
        u = int(input[ptr]); ptr +=1
        v = int(input[ptr]); ptr +=1
        adj[u].append(v)
        adj[v].append(u)
    
    s = []
    for _ in range(N):
        s_j = int(input[ptr]); ptr +=1
        s.append(s_j)
    
    x = list(map(int, input[ptr:ptr+K]))
    ptr += K
    
    infected = set()
    immune = set()
    quarantined = set()
    recovery_map = defaultdict(list)
    spread_map = defaultdict(list)
    
    # Initialize infected rooms
    for room in x:
        if room in immune or room in quarantined:
            continue
        infected.add(room)
        spread_day = 0 + s[room-1]
        spread_map[spread_day].append(room)
        recovery_day = 0 + P_val
        recovery_map[recovery_day].append(room)
    
    quarantine_count = 0
    
    day = 0
    while True:
        # Step 1: Recover rooms on current day
        if day in recovery_map:
            recoveries = recovery_map[day]
            for room in recoveries:
                if room in infected:
                    infected.remove(room)
                    immune.add(room)
            del recovery_map[day]
        
        # Step 2: Process spreads on current day
        if day in spread_map:
            spreads = spread_map[day]
            new_infections = dict()  # key: neighbor, value: infection day
            temp_quarantine = set()
            for room in spreads:
                if room in quarantined:
                    continue
                for neighbor in adj[room]:
                    if neighbor in immune or neighbor in quarantined:
                        continue
                    if neighbor not in infected and neighbor not in new_infections:
                        new_infections[neighbor] = day
                    else:
                        temp_quarantine.add(neighbor)
            # Process new_infections
            for neighbor in new_infections:
                inf_day = new_infections[neighbor]
                if neighbor not in infected:
                    infected.add(neighbor)
                    # Schedule spread and recovery events
                    spread_day_neighbor = inf_day + s[neighbor-1]
                    spread_map[spread_day_neighbor].append(neighbor)
                    recovery_day_neighbor = inf_day + P_val
                    recovery_map[recovery_day_neighbor].append(neighbor)
                else:
                    # Neighbor was infected before this day's new_infections
                    temp_quarantine.add(neighbor)
            # Step 3: Check quarantines
            for q_room in temp_quarantine:
                if q_room in infected and q_room not in quarantined:
                    quarantined.add(q_room)
                    infected.remove(q_room)
                    quarantine_count += 1
            del spread_map[day]
        
        # Move to next day if infected exists
        day += 1
        if not infected:
            break
    
    print(quarantine_count)

if __name__ == '__main__':
    main()
0