結果

問題 No.901 K-ary εxtrεεmε
ユーザー terasaterasa
提出日時 2023-01-15 17:13:10
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 641 ms / 3,000 ms
コード長 4,628 bytes
コンパイル時間 1,422 ms
コンパイル使用メモリ 86,996 KB
実行使用メモリ 124,972 KB
最終ジャッジ日時 2023-08-28 09:23:09
合計ジャッジ時間 20,497 ms
ジャッジサーバーID
(参考情報)
judge14 / judge12
このコードへのチャレンジ
(要ログイン)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 413 ms
124,972 KB
testcase_01 AC 162 ms
80,628 KB
testcase_02 AC 194 ms
83,368 KB
testcase_03 AC 191 ms
83,076 KB
testcase_04 AC 192 ms
83,564 KB
testcase_05 AC 185 ms
82,984 KB
testcase_06 AC 189 ms
82,932 KB
testcase_07 AC 599 ms
121,220 KB
testcase_08 AC 597 ms
120,512 KB
testcase_09 AC 602 ms
120,492 KB
testcase_10 AC 607 ms
120,524 KB
testcase_11 AC 615 ms
121,136 KB
testcase_12 AC 623 ms
120,472 KB
testcase_13 AC 628 ms
120,412 KB
testcase_14 AC 612 ms
120,396 KB
testcase_15 AC 597 ms
120,468 KB
testcase_16 AC 599 ms
120,632 KB
testcase_17 AC 631 ms
121,176 KB
testcase_18 AC 628 ms
120,660 KB
testcase_19 AC 641 ms
121,312 KB
testcase_20 AC 634 ms
121,160 KB
testcase_21 AC 631 ms
120,756 KB
testcase_22 AC 620 ms
121,320 KB
testcase_23 AC 611 ms
120,488 KB
testcase_24 AC 615 ms
120,876 KB
testcase_25 AC 609 ms
120,648 KB
testcase_26 AC 604 ms
120,944 KB
testcase_27 AC 591 ms
120,488 KB
testcase_28 AC 595 ms
120,248 KB
testcase_29 AC 589 ms
120,684 KB
権限があれば一括ダウンロードができます

ソースコード

diff #

from typing import List, Tuple, Callable, TypeVar, Optional
import sys
import itertools
import heapq
import bisect
import math
from collections import deque, defaultdict
from functools import lru_cache, cmp_to_key

input = sys.stdin.readline

if __file__ != 'prog.py':
    sys.setrecursionlimit(10 ** 6)


def readints(): return map(int, input().split())
def readlist(): return list(readints())
def readstr(): return input()[:-1]
def readlist1(): return list(map(lambda x: int(x) - 1, input().split()))


class CumulativeSum:
    def __init__(self, A):
        self.S = [0]
        acc = 0
        for a in A:
            acc += a
            self.S.append(acc)

    def get(self, l, r):
        """return sum(A[l:r]), i.e. sum of A[x] (l <= x < r) """
        return self.S[r] - self.S[l]


class HLD:
    # reference: https://codeforces.com/blog/entry/53170
    def __init__(self, N, E, root: int = 0):
        self.N = N
        self.E = E
        self.root = root

        self.D = [0] * self.N
        self.par = [-1] * self.N
        self.sz = [0] * self.N
        self.top = [0] * self.N

        self.ord = [None] * self.N

        self._dfs_sz()
        self._dfs_hld()

    def path_query_range(self, u: int, v: int, is_edge_query: bool = False) -> List[Tuple[int, int]]:
        """return list of [l, r) ranges that cover u-v path"""
        ret = []
        while True:
            if self.ord[u] > self.ord[v]:
                u, v = v, u
            if self.top[u] == self.top[v]:
                ret.append((self.ord[u] + is_edge_query, self.ord[v] + 1))
                return ret
            ret.append((self.ord[self.top[v]], self.ord[v] + 1))
            v = self.par[self.top[v]]

    def subtree_query_range(self, v: int, is_edge_query: bool = False) -> Tuple[int, int]:
        """return [l, r) range that cover vertices of subtree v"""
        return (self.ord[v] + is_edge_query, self.ord[v] + self.sz[v])

    def get_index(self, v: int) -> int:
        """return euler tour order of given vertex"""
        return self.ord[v]

    def lca(self, u, v):
        while True:
            if self.ord[u] > self.ord[v]:
                u, v = v, u
            if self.top[u] == self.top[v]:
                return u
            v = self.par[self.top[v]]

    def dist(self, u, v):
        return self.D[u] + self.D[v] - 2 * self.D[self.lca(u, v)]

    def _dfs_sz(self):
        stack = [(self.root, -1)]
        while stack:
            v, p = stack.pop()
            if v < 0:
                v = ~v
                self.sz[v] = 1
                for i, dst in enumerate(self.E[v]):
                    if dst == p:
                        continue
                    self.sz[v] += self.sz[dst]
                    # v -> E[v][0] will be heavy path
                    if self.sz[self.E[v][0]] < self.sz[dst]:
                        self.E[v][0], self.E[v][i] = self.E[v][i], self.E[v][0]
            else:
                if ~p:
                    self.D[v] = self.D[p] + 1
                    self.par[v] = p
                # avoid first element of E[v] is parent of vertex v if v has some children
                if len(self.E[v]) >= 2 and self.E[v][0] == p:
                    self.E[v][0], self.E[v][1] = self.E[v][1], self.E[v][0]
                stack.append((~v, p))
                for dst in self.E[v]:
                    if dst == p:
                        continue
                    stack.append((dst, v))

    def _dfs_hld(self):
        stack = [(self.root, -1)]
        cnt = 0
        while stack:
            v, p = stack.pop()
            self.ord[v] = cnt
            cnt += 1
            heavy_path_idx = len(self.E[v]) - 1
            for i, dst in enumerate(self.E[v][::-1]):
                if dst == p:
                    continue
                # top[dst] is top[v] if v -> dst is heavy path otherwise dst itself
                self.top[dst] = self.top[v] if i == heavy_path_idx else dst
                stack.append((dst, v))


N = int(input())
E = [[] for _ in range(N)]
edges = []
for _ in range(N - 1):
    u, v, w = readints()
    E[u].append(v)
    E[v].append(u)
    edges.append((u, v, w))
solver = HLD(N, E)
W = [0] * N
for u, v, w in edges:
    if solver.ord[u] > solver.ord[v]:
        u, v = v, u
    W[solver.get_index(v)] = w
S = CumulativeSum(W)
Q = int(input())
for _ in range(Q):
    k, *X = readlist()
    X = sorted(X, key=lambda x: solver.ord[x])
    ans = 0
    for i in range(k):
        s = X[i]
        t = X[(i + 1) % k]
        for l, r in solver.path_query_range(s, t, is_edge_query=True):
            ans += S.get(l, r)
    print(ans // 2)
0