結果

問題 No.1197 モンスターショー
ユーザー gew1fw
提出日時 2025-06-12 15:36:03
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 3,571 bytes
コンパイル時間 202 ms
コンパイル使用メモリ 82,176 KB
実行使用メモリ 117,120 KB
最終ジャッジ日時 2025-06-12 15:36:34
合計ジャッジ時間 6,188 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample -- * 2
other AC * 7 TLE * 1 -- * 33
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
from collections import deque

def main():
    sys.setrecursionlimit(1 << 25)
    input = sys.stdin.read().split()
    ptr = 0
    N = int(input[ptr])
    ptr += 1
    K = int(input[ptr])
    ptr += 1
    Q = int(input[ptr])
    ptr += 1

    C = list(map(int, input[ptr:ptr+K]))
    ptr += K

    edges = [[] for _ in range(N+1)]
    for _ in range(N-1):
        a = int(input[ptr])
        ptr += 1
        b = int(input[ptr])
        ptr += 1
        edges[a].append(b)
        edges[b].append(a)

    root = 1
    depth = [0] * (N + 1)
    parent = [None] * (N + 1)
    visited = [False] * (N + 1)
    q = deque()
    q.append(root)
    visited[root] = True
    while q:
        u = q.popleft()
        for v in edges[u]:
            if not visited[v] and v != parent[u]:
                parent[v] = u
                depth[v] = depth[u] + 1
                visited[v] = True
                q.append(v)

    in_time = [0] * (N + 1)
    out_time = [0] * (N + 1)
    time = 1
    stack = [(root, False)]
    while stack:
        node, processed = stack.pop()
        if processed:
            out_time[node] = time
            time += 1
        else:
            in_time[node] = time
            time += 1
            stack.append((node, True))
            for v in reversed(edges[node]):
                if v != parent[node]:
                    stack.append((v, False))

    fenwick_size = N * 2
    fenwick = [0] * (fenwick_size + 2)

    def fenwick_update(idx, delta):
        idx += 1
        while idx <= fenwick_size + 1:
            fenwick[idx] += delta
            idx += idx & -idx

    def fenwick_query(idx):
        idx += 1
        res = 0
        while idx > 0:
            res += fenwick[idx]
            idx -= idx & -idx
        return res

    def range_query(l, r):
        if l > r:
            return 0
        return fenwick_query(r) - fenwick_query(l - 1)

    sum_depth = 0
    for c in C:
        u = c
        fenwick_update(in_time[u], 1)
        sum_depth += depth[u]

    for _ in range(Q):
        query = input[ptr]
        ptr += 1
        if query == '1':
            p = int(input[ptr])
            ptr += 1
            d = int(input[ptr])
            ptr += 1
            p -= 1
            old = C[p]
            new = d
            sum_depth -= depth[old]
            sum_depth += depth[new]
            fenwick_update(in_time[old], -1)
            fenwick_update(in_time[new], 1)
            C[p] = new
        elif query == '2':
            e = int(input[ptr])
            ptr += 1
            term1 = K * depth[e]
            term2 = sum_depth
            term3 = 0
            x = e
            is_root = (x == root)
            while True:
                if x is None:
                    break
                current_in = in_time[x]
                current_out = out_time[x]
                count_sub = range_query(current_in, current_out)
                if x == root:
                    count_parent = 0
                    if parent[x] is not None:
                        parent_in = in_time[parent[x]]
                        parent_out = out_time[parent[x]]
                        count_parent = range_query(parent_in, parent_out)
                    count_LCA = count_sub - count_parent
                else:
                    count_LCA = count_sub
                term3 += count_LCA * depth[x]
                if parent[x] is None:
                    break
                x = parent[x]
            total = term1 + term2 - 2 * term3
            print(total)

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