結果
問題 | No.386 貪欲な領主 |
ユーザー |
![]() |
提出日時 | 2023-01-15 16:26:06 |
言語 | PyPy3 (7.3.15) |
結果 |
AC
|
実行時間 | 597 ms / 2,000 ms |
コード長 | 4,400 bytes |
コンパイル時間 | 435 ms |
コンパイル使用メモリ | 82,324 KB |
実行使用メモリ | 107,724 KB |
最終ジャッジ日時 | 2024-12-29 04:20:12 |
合計ジャッジ時間 | 5,585 ms |
ジャッジサーバーID (参考情報) |
judge1 / judge4 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
other | AC * 16 |
ソースコード
from typing import List, Tuple, Callable, TypeVar, Optionalimport sysimport itertoolsimport heapqimport bisectimport mathfrom collections import deque, defaultdictfrom functools import lru_cache, cmp_to_keyinput = sys.stdin.readlineif __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 = 0for a in A:acc += aself.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/53170def __init__(self, N, E, root: int = 0):self.N = Nself.E = Eself.root = rootself.D = [0] * self.Nself.par = [-1] * self.Nself.sz = [0] * self.Nself.top = [0] * self.Nself.ord = [None] * self.Nself._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, uif self.top[u] == self.top[v]:ret.append((self.ord[u] + is_edge_query, self.ord[v] + 1))return retret.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, uif self.top[u] == self.top[v]:return uv = 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 = ~vself.sz[v] = 1for i, dst in enumerate(self.E[v]):if dst == p:continueself.sz[v] += self.sz[dst]# v -> E[v][0] will be heavy pathif 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] + 1self.par[v] = p# avoid first element of E[v] is parent of vertex v if v has some childrenif 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:continuestack.append((dst, v))def _dfs_hld(self):stack = [(self.root, -1)]cnt = 0while stack:v, p = stack.pop()self.ord[v] = cntcnt += 1heavy_path_idx = len(self.E[v]) - 1for 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 itselfself.top[dst] = self.top[v] if i == heavy_path_idx else dststack.append((dst, v))N = int(input())E = [[] for _ in range(N)]for _ in range(N - 1):a, b = readints()E[a].append(b)E[b].append(a)solver = HLD(N, E)A = [None] * Nfor i in range(N):u = int(input())A[solver.get_index(i)] = uS = CumulativeSum(A)Q = int(input())ans = 0for _ in range(Q):a, b, c = readints()for l, r in solver.path_query_range(a, b):ans += S.get(l, r) * cprint(ans)