#1024361 (PyPy3) No.901 K-ary εxtrεεmε

提出ソース
結果

問題	No.901 K-ary εxtrεεmε
ユーザー	Navier_Boltzmann
提出日時	2024-11-04 21:21:59
言語	PyPy3 (7.3.15)
結果	WA
実行時間	-
コード長	11,277 bytes
コンパイル時間	328 ms
コンパイル使用メモリ	82,324 KB
実行使用メモリ	221,200 KB
最終ジャッジ日時	2024-11-04 21:22:54
合計ジャッジ時間	48,079 ms
ジャッジサーバーID （参考情報）	judge4 / judge1
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	WA	-
testcase_01	WA	-
testcase_02	WA	-
testcase_03	WA	-
testcase_04	WA	-
testcase_05	WA	-
testcase_06	WA	-
testcase_07	WA	-
testcase_08	WA	-
testcase_09	WA	-
testcase_10	WA	-
testcase_11	WA	-
testcase_12	WA	-
testcase_13	WA	-
testcase_14	WA	-
testcase_15	WA	-
testcase_16	WA	-
testcase_17	WA	-
testcase_18	WA	-
testcase_19	WA	-
testcase_20	WA	-
testcase_21	WA	-
testcase_22	WA	-
testcase_23	WA	-
testcase_24	WA	-
testcase_25	WA	-
testcase_26	WA	-
testcase_27	AC	1,090 ms 154,592 KB
testcase_28	AC	1,081 ms 153,964 KB
testcase_29	AC	1,072 ms 153,068 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
# import pypyjit

# pypyjit.set_param("max_unroll_recursion=-1")
import sys
from collections import *
from functools import *

# sys.setrecursionlimit(10**6)


class SegTree:
    """
    Segment Tree
    """

    def __init__(self, init_val, segfunc, ide_ele):
        """
        初期化

        init_val: 配列の初期値
        """
        n = len(init_val)
        self.segfunc = segfunc
        self.ide_ele = ide_ele
        self.num = 1 << (n - 1).bit_length()
        self.tree = [ide_ele] * 2 * self.num
        # 配列の値を葉にセット
        for i in range(n):
            self.tree[self.num + i] = init_val[i]
        # 構築していく
        for i in range(self.num - 1, 0, -1):
            self.tree[i] = segfunc(self.tree[2 * i], self.tree[2 * i + 1])

    def update(self, k, x):
        """
        k番目の値をxに更新

        k: index(0-index)
        x: update value
        """
        k += self.num
        self.tree[k] = x
        while k > 1:
            self.tree[k >> 1] = self.segfunc(self.tree[k], self.tree[k ^ 1])
            k >>= 1

    def query(self, l, r):
        """
        [l, r)のsegfuncしたものを得る

        l: index(0-index)
        r: index(0-index)
        """
        res = self.ide_ele

        l += self.num
        r += self.num
        while l < r:
            if l & 1:
                res = self.segfunc(res, self.tree[l])
                l += 1
            if r & 1:
                res = self.segfunc(res, self.tree[r - 1])
            l >>= 1
            r >>= 1
        return res


class HLD:

    ### HL分解をしてIDを振りなおしたものに対して、パスに含まれる区間を返す
    ### SegTreeにのせる配列はIDを並び替えたもの

    def __init__(self, e, root=0):

        self.N = len(e)
        self.e = e
        par = [-1] * self.N
        sub = [-1] * self.N
        self.root = root
        dist = [-1] * self.N
        v = deque()
        dist[root] = 0
        v.append(root)
        while v:
            x = v.popleft()
            for ix in e[x]:
                if dist[ix] != -1:
                    continue
                dist[ix] = dist[x] + 1
                v.append(ix)

        H = [(-dist[i], i) for i in range(self.N)]
        H.sort()
        for h, i in H:
            tmp = 1
            for ix in e[i]:
                if sub[ix] == -1:
                    par[i] = ix
                else:
                    tmp += sub[ix]
            sub[i] = tmp

        self.ID = [-1] * self.N
        self.ID[self.root] = 0
        self.HEAD = [-1] * self.N
        head = [-1] * self.N
        self.PAR = [-1] * self.N
        visited = [False] * self.N
        self.HEAD[0] = 0
        head[self.root] = 0
        depth = [-1] * self.N
        depth[self.root] = 0
        self.DEPTH = [-1] * self.N
        self.DEPTH[0] = 0
        cnt = 0
        v = deque([self.root])
        self.SUB = [0] * self.N
        self.SUB[0] = self.N
        while v:
            x = v.popleft()
            visited[x] = True
            self.ID[x] = cnt
            cnt += 1
            n = len(self.e[x])
            tmp = [(sub[ix], ix) for ix in self.e[x]]
            tmp.sort()
            flg = 0
            if x == self.root:
                flg -= 1
            for _, ix in tmp:
                flg += 1
                if visited[ix]:
                    continue
                v.appendleft(ix)
                if flg == n - 1:
                    head[ix] = head[x]
                    depth[ix] = depth[x]
                else:
                    head[ix] = ix
                    depth[ix] = depth[x] + 1

        for i in range(self.N):
            self.PAR[self.ID[i]] = self.ID[par[i]]
            self.HEAD[self.ID[i]] = self.ID[head[i]]
            self.DEPTH[self.ID[i]] = depth[i]
            self.SUB[self.ID[i]] = sub[i]

    def path_query(self, l, r):
        L = self.ID[l]
        R = self.ID[r]
        res = []
        if self.DEPTH[L] < self.DEPTH[R]:
            L, R = R, L

        while self.DEPTH[L] != self.DEPTH[R]:
            tmp = (self.HEAD[L], L + 1)
            res.append(tmp)
            L = self.PAR[self.HEAD[L]]

        while self.HEAD[L] != self.HEAD[R]:
            tmp = (self.HEAD[L], L + 1)
            res.append(tmp)
            L = self.PAR[self.HEAD[L]]
            tmp = (self.HEAD[R], R + 1)
            res.append(tmp)
            R = self.PAR[self.HEAD[R]]

        if L > R:
            L, R = R, L

        tmp = (L, R + 1)
        res.append(tmp)

        return res

    def sub_query(self, k):

        K = self.ID[k]

        return (K, K + self.SUB[K])


class HLD_SegTree:

    def __init__(self, e, init_val, segfunc, ide_ele, root=0):

        self.hld = HLD(e, root=root)
        self.ID = self.hld.ID[:]
        self.N = len(e)
        A = [0] * self.N
        for i, idx in enumerate(self.ID):
            A[idx] = init_val[i]
        self.seg = SegTree(A, segfunc, ide_ele)
        self.segfunc = segfunc
        self.ide_ele = ide_ele

    def path_query(self, l, r):
        res = self.ide_ele
        for _l, _r in self.hld.path_query(l, r):
            res = self.segfunc(res, self.seg.query(_l, _r))
        return res

    def sub_query(self, l, r):

        _l, _r = self.hld.sub_query(l, r)
        return self.seg.query(_l, _r)


class AuxiliaryTree:
    def __init__(self, e):
        """
        Initializes the AuxiliaryTree with an adjacency list.

        :param e: List of lists representing the adjacency list of the tree.
        """
        self.n = len(e)
        self.adj = e
        self.depth = [0] * self.n
        self.parent = [-1] * self.n
        self.euler = []
        self.first_occurrence = [-1] * self.n
        self.log = [0] * (2 * self.n)
        self.sparse_table = []

        # Initialize log array
        self._initialize_log()

        # Perform DFS to populate Euler tour
        self._iterative_dfs(0, -1, 0)

        # Build the sparse table
        self._build_sparse_table()

    def _initialize_log(self):
        """
        Initializes the log array for use in the sparse table.
        """
        for i in range(2, len(self.log)):
            self.log[i] = self.log[i // 2] + 1

    def _iterative_dfs(self, start_node, par, dep):
        """
        Performs an iterative depth-first search (DFS) to populate the Euler tour,
        depth, and parent information for the tree.

        :param start_node: Starting node for DFS.
        :param par: Parent of the starting node.
        :param dep: Depth of the starting node.
        """
        stack = [(start_node, par, dep)]
        while stack:
            node, par, dep = stack.pop()
            if self.first_occurrence[node] == -1:
                self.first_occurrence[node] = len(self.euler)
            self.euler.append(node)
            self.parent[node] = par
            self.depth[node] = dep

            for neighbor in reversed(self.adj[node]):
                if neighbor != par:
                    stack.append((neighbor, node, dep + 1))
                    self.euler.append(node)

    def _build_sparse_table(self):
        """
        Builds a sparse table for range minimum queries on the Euler tour,
        which is used to efficiently compute the lowest common ancestor (LCA).
        """
        m = len(self.euler)
        max_log = self.log[m] + 1
        self.sparse_table = [[0] * m for _ in range(max_log)]
        for i in range(m):
            self.sparse_table[0][i] = self.euler[i]

        j = 1
        while (1 << j) <= m:
            i = 0
            while (i + (1 << j) - 1) < m:
                if (
                    self.depth[self.sparse_table[j - 1][i]]
                    < self.depth[self.sparse_table[j - 1][i + (1 << (j - 1))]]
                ):
                    self.sparse_table[j][i] = self.sparse_table[j - 1][i]
                else:
                    self.sparse_table[j][i] = self.sparse_table[j - 1][
                        i + (1 << (j - 1))
                    ]
                i += 1
            j += 1

    def lca(self, u, v):
        """
        Computes the lowest common ancestor (LCA) of two nodes u and v.

        :param u: First node.
        :param v: Second node.
        :return: The LCA of nodes u and v.
        """
        left = self.first_occurrence[u]
        right = self.first_occurrence[v]
        if left > right:
            left, right = right, left
        length = right - left + 1
        j = self.log[length]

        # Compare depths of the two intervals in the sparse table
        left_interval = self.sparse_table[j][left]
        right_interval = self.sparse_table[j][right - (1 << j) + 1]

        if self.depth[left_interval] < self.depth[right_interval]:
            return left_interval
        else:
            return right_interval

    def build_auxiliary_tree(self, nodes):
        """
        Constructs an auxiliary tree from a subset of nodes.

        :param nodes: List of nodes to include in the auxiliary tree.
        :return: A dictionary representing the adjacency list of the auxiliary tree.
        """
        nodes = sorted(nodes, key=lambda x: self.first_occurrence[x])
        stack = []
        auxiliary_tree = defaultdict(set)  # Use a set to avoid duplicates

        for node in nodes:
            if not stack:
                stack.append(node)
                continue

            lca_node = self.lca(stack[-1], node)
            while (
                len(stack) > 1
                and self.first_occurrence[stack[-2]] >= self.first_occurrence[lca_node]
            ):
                top = stack.pop()
                auxiliary_tree[stack[-1]].add(top)
                auxiliary_tree[top].add(stack[-1])  # Add reverse edge

            if stack[-1] != lca_node:
                top = stack.pop()
                auxiliary_tree[lca_node].add(top)
                auxiliary_tree[top].add(lca_node)  # Add reverse edge
                stack.append(lca_node)

            stack.append(node)

        # Ensure all remaining nodes in the stack are added to the auxiliary tree
        while len(stack) > 1:
            top = stack.pop()
            auxiliary_tree[stack[-1]].add(top)
            auxiliary_tree[top].add(stack[-1])  # Add reverse edge

        # Convert sets back to lists for the final output
        return {k: list(v) for k, v in auxiliary_tree.items()}


N = int(input())
e = [[] for _ in range(N)]
W = [0] * N
edge = defaultdict(int)
for _ in range(N - 1):
    u, v, w = map(int, input().split())
    e[u].append(v)
    e[v].append(u)
    edge[(u, v)] = w
tree = AuxiliaryTree(e)
for (u, v), w in edge.items():
    if v == tree.parent[u]:
        W[u] = w
    else:
        W[v] = w
hld = HLD_SegTree(e, W, lambda x, y: x + y, 0)
Q = int(input())
for _ in range(Q):
    _, *K = map(int, input().split())
    # K = [k - 1 for k in K]
    at = tree.build_auxiliary_tree(K)
    ans = 0
    for i, v in at.items():
        for j in v:
            ans += hld.path_query(i, j) - W[tree.lca(i, j)]
    #         print(i, j, tree.lca(i, j), hld.path_query(i, j), W[tree.lca(i, j)])
    # print(at, K)
    print(ans // 2)
# for i in range(N):
#     for j in range(i + 1, N):
#         print(i, j, tree.lca(i, j))
yukicoder

結果

テストケース

ソースコード
