結果
問題 | No.386 貪欲な領主 |
ユーザー |
|
提出日時 | 2021-09-16 02:54:38 |
言語 | PyPy3 (7.3.15) |
結果 |
AC
|
実行時間 | 618 ms / 2,000 ms |
コード長 | 3,884 bytes |
コンパイル時間 | 179 ms |
コンパイル使用メモリ | 82,236 KB |
実行使用メモリ | 111,124 KB |
最終ジャッジ日時 | 2024-06-29 03:10:37 |
合計ジャッジ時間 | 6,004 ms |
ジャッジサーバーID (参考情報) |
judge3 / judge2 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
other | AC * 16 |
ソースコード
#!/usr/bin/env python3import sysclass LcaDoubling:# 木であれば任意の点を根と見做せる。def __init__(self, N, root=0):self.N = Nself.root = rootself.G = [[] for _ in range(N)]self.depths = [-1] * Nself.distances = [-1] * Nself.ancestors = []returndef addEdge(self, a: int, b: int, cost: int):self.G[a].append((cost, b))self.G[b].append((cost, a))returndef build(self):prevAncestors = self._dfs()self.ancestors.append(prevAncestors)d = 1max_depth = max(self.depths)while d < max_depth:nextAncestors = [prevAncestors[p] for p in prevAncestors]self.ancestors.append(nextAncestors)d <<= 1prevAncestors = nextAncestorsreturndef _dfs(self):q = [(self.root, -1, 0, 0)]directAncestors = [-1] * (self.N + 1) # 頂点数より1個長くし、存在しないことを-1で表す。末尾(-1)要素は常に-1while q:now, parent, dep, dist = q.pop()directAncestors[now] = parentself.depths[now] = depself.distances[now] = distfor cost, next in self.G[now]:if next != parent:q.append((next, now, dep + 1, dist + cost))return directAncestorsdef getLca(self, nodeA: int, nodeB: int):depthA, depthB = self.depths[nodeA], self.depths[nodeB]if depthA > depthB:nodeA, nodeB = nodeB, nodeAdepthA, depthB = depthB, depthA# 2ノードを同じ深さまで揃える。tu = nodeAtv = self.upstream(nodeB, depthB - depthA)# 遡上させて行き2つが衝突する位置が共通祖先。if nodeA == tv:return nodeAfor k in range(depthA.bit_length() - 1, -1, -1):mu = self.ancestors[k][tu]mv = self.ancestors[k][tv]if mu != mv:tu = mutv = mvlca = self.ancestors[0][tu]assert lca == self.ancestors[0][tv]return lca# 2つのノードの間の距離を返す。def getDistance(self, nodeA, nodeB):lca = self.getLca(nodeA, nodeB)return self.distances[nodeA] + self.distances[nodeB] - 2 * self.distances[lca]# targetNodeが2つのノード間のパス上に存在するかを返す。def isOnPath(self, nodeA: int, nodeB: int, evalNode: int):return self.getDistance(nodeA, nodeB) == self.getDistance(nodeA, evalNode) + self.getDistance(evalNode, nodeB)# ノードvからk個遡上したノードを返す。def upstream(self, v, k):i = 0while k:if k & 1:v = self.ancestors[i][v]k >>= 1i += 1return vdef main():N = int(input())ld = LcaDoubling(N)for _ in range(N - 1):a, b = map(int, input().split())ld.addEdge(a, b, 1)ld.build()U = [int(input()) for _ in range(N)]from collections import dequedef bfs(edges: "List[to]", start_node: int) -> list:q = deque()dist = [0] * len(edges)q.append(start_node)dist[start_node] = U[0]while q:now = q.popleft()for _, next in edges[now]:if dist[next] != 0:continueq.append(next)dist[next] = dist[now] + U[next]return distd = bfs(ld.G, 0)M = int(input())ans = 0for _ in range(M):u, v, e = map(int, input().split())lca = ld.getLca(u, v)ans += (d[u] + d[v] - 2 * d[lca] + U[lca]) * eprint(ans)returnif __name__ == '__main__':main()