#1178507 (PyPy3) No.3222 Let the World Forget Me

• 問題: No.3222 Let the World Forget Me
• ユーザー: LyricalMaestro
• 提出日時: 2026-08-23 03:22:44
• 言語: PyPy3 (7.3.17)
• 結果: AC
• 実行時間: 430 ms / 2,000 ms
• コード長: 3,583 bytes
• コンパイル時間: 251 ms
• コンパイル使用メモリ: 95,340 KB
• 実行使用メモリ: 149,436 KB
• 最終ジャッジ日時: 2026-08-23 03:22:59
• 合計ジャッジ時間: 14,378 ms
• ジャッジサーバーID （参考情報）: judge1_0 / judge3_1

ファイルパターン	結果
sample	AC * 1
other	AC * 31

ソースコード

from collections import deque

MAX_INT = 10 ** 18

class SegmentTree:
    """
    非再帰版セグメント木。
    更新は「加法」、取得は「最大値」のもの限定。
    """

    def __init__(self, init_array):
        n = 1
        while n < len(init_array):
            n *= 2
        
        self.size = n
        self.array = [[-MAX_INT, -MAX_INT] for _ in range(2 * self.size)]
        for i, a in enumerate(init_array):
            self.array[self.size + i][0] = a
            self.array[self.size + i][1] = i
        
        end_index = self.size
        start_index = end_index // 2
        while start_index >= 1:
            for i in range(start_index, end_index):
                self._op(self.array[i], self.array[2 * i], self.array[2 * i + 1])
            end_index = start_index
            start_index = end_index // 2

    def _op(self, array, left, right):
        if left[0] > right[0]:
            array[0] = left[0]
            array[1] = left[1]
        elif left[0] < right[0]:
            array[0] = right[0]
            array[1] = right[1]
        else:
            array[0] = left[0]
            array[1] = min(left[1], right[1])


    def set(self, x, a):
        index = self.size + x
        self.array[index][0] = a
        self.array[index][1] = x
        while index > 1:
            index //= 2
            self._op(self.array[index], self.array[2 * index], self.array[2 * index+ 1])

    def get_max(self, l, r):
        L = self.size + l; R = self.size + r

        # 2. 区間[l, r)の最大値を求める
        s = [- MAX_INT, - MAX_INT]
        while L < R:
            if R & 1:
                R -= 1
                self._op(s, s, self.array[R])
            if L & 1:
                self._op(s, s, self.array[L])
                L += 1
            L >>= 1; R >>= 1
        return s


def main():
    N, M = map(int, input().split())
    P = list(map(int, input().split()))
    next_nodes = [set() for _ in range(N)]
    for _ in range(N - 1):
        a, b = map(int, input().split())
        next_nodes[a - 1].add(b - 1)
        next_nodes[b - 1].add(a - 1)
    C = list(map(int, input().split()))

    # このまま放置していたら感染されるような場合、いつnodeが感染されるか計算
    dists = [MAX_INT] * N
    queue = deque()
    for c in C:
        dists[c - 1] = 0
        queue.append(c -  1)

    while len(queue) > 0:
        v = queue.popleft()
        for w in next_nodes[v]:
            if dists[w] == MAX_INT:
                dists[w] = dists[v] + 1
                queue.append(w)

    dists_map = {}
    for i in range(N):
        d = dists[i]
        if d not in dists_map:
            dists_map[d] = []
        dists_map[d].append(i)

    seg_tree = SegmentTree([0] * N)
    for i in range(N):
        if len(next_nodes[i]) == 1:
            seg_tree.set(i, P[i])

    answer = 0
    enable = [True] * N
    for t in range(N + 1):
        if t in dists_map:
            for index in dists_map[t]:
                enable[index] = False
                seg_tree.set(index, 0)

        s = seg_tree.get_max(0, N)
        if s[0] >= 1:
            index = s[1]
            answer += s[0]
            seg_tree.set(index, 0)

            for w in next_nodes[index]:
                next_nodes[w].remove(index)
                if len(next_nodes[w]) == 1 and enable[w]:
                    seg_tree.set(w, P[w])
            next_nodes[index].clear()
    print(answer)



        

        


            



            
        



        


if __name__ == "__main__":
    main()