結果

問題 No.899 γatheree
ユーザー lam6er
提出日時 2025-04-09 20:56:08
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 1,798 bytes
コンパイル時間 571 ms
コンパイル使用メモリ 82,444 KB
実行使用メモリ 127,960 KB
最終ジャッジ日時 2025-04-09 20:56:39
合計ジャッジ時間 12,807 ms
ジャッジサーバーID
(参考情報)
judge3 / judge2
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 1
other AC * 20 TLE * 1 -- * 2
権限があれば一括ダウンロードができます

ソースコード

diff #

import sys
from collections import deque

def main():
    input = sys.stdin.read
    data = input().split()
    ptr = 0
    N = int(data[ptr])
    ptr += 1
    edges = [[] for _ in range(N)]
    for _ in range(N - 1):
        u = int(data[ptr])
        v = int(data[ptr + 1])
        edges[u].append(v)
        edges[v].append(u)
        ptr += 2
    A = list(map(int, data[ptr:ptr + N]))
    ptr += N
    Q = int(data[ptr])
    ptr += 1
    queries = list(map(int, data[ptr:ptr + Q]))
    
    # Precompute adjacency lists
    adj = [[] for _ in range(N)]
    for u in range(N):
        for v in edges[u]:
            if v not in adj[u]:
                adj[u].append(v)
    
    # Precompute sum_adj
    sum_adj = [0] * N
    for u in range(N):
        sum_adj[u] = sum(A[v] for v in adj[u])
    
    # To avoid recomputing, we'll maintain A and sum_adj dynamically
    # We'll handle each query by iterating through adjacency lists
    for x in queries:
        # Collect all nodes in S(x)
        s = {x}
        for v in adj[x]:
            s.add(v)
            for u in adj[v]:
                s.add(u)
        
        # Compute sum_total
        sum_total = 0
        cleared = []
        for y in s:
            sum_total += A[y]
        # Set all y in s (except x) to 0, and update sum_adj
        for y in s:
            if y == x:
                continue
            if A[y] != 0:
                val = A[y]
                A[y] = 0
                for u in adj[y]:
                    sum_adj[u] -= val
        # Set x's A to sum_total and update sum_adj for x's neighbors
        old_Ax = A[x]
        A[x] = sum_total
        delta = sum_total - old_Ax
        for v in adj[x]:
            sum_adj[v] += delta
        
        print(sum_total)

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