結果
問題 | No.1098 LCAs |
ユーザー | None |
提出日時 | 2021-02-26 00:27:06 |
言語 | PyPy3 (7.3.15) |
結果 |
AC
|
実行時間 | 944 ms / 2,000 ms |
コード長 | 8,528 bytes |
コンパイル時間 | 182 ms |
コンパイル使用メモリ | 82,048 KB |
実行使用メモリ | 174,488 KB |
最終ジャッジ日時 | 2024-10-01 15:09:12 |
合計ジャッジ時間 | 14,329 ms |
ジャッジサーバーID (参考情報) |
judge1 / judge2 |
(要ログイン)
テストケース
テストケース表示入力 | 結果 | 実行時間 実行使用メモリ |
---|---|---|
testcase_00 | AC | 42 ms
53,632 KB |
testcase_01 | AC | 43 ms
53,632 KB |
testcase_02 | AC | 42 ms
53,632 KB |
testcase_03 | AC | 43 ms
53,760 KB |
testcase_04 | AC | 41 ms
53,120 KB |
testcase_05 | AC | 42 ms
53,888 KB |
testcase_06 | AC | 42 ms
53,632 KB |
testcase_07 | AC | 41 ms
53,248 KB |
testcase_08 | AC | 42 ms
53,504 KB |
testcase_09 | AC | 41 ms
53,376 KB |
testcase_10 | AC | 40 ms
53,120 KB |
testcase_11 | AC | 41 ms
53,760 KB |
testcase_12 | AC | 42 ms
53,376 KB |
testcase_13 | AC | 64 ms
68,736 KB |
testcase_14 | AC | 63 ms
68,096 KB |
testcase_15 | AC | 66 ms
68,480 KB |
testcase_16 | AC | 63 ms
68,352 KB |
testcase_17 | AC | 65 ms
68,096 KB |
testcase_18 | AC | 890 ms
152,064 KB |
testcase_19 | AC | 901 ms
151,808 KB |
testcase_20 | AC | 918 ms
151,864 KB |
testcase_21 | AC | 910 ms
151,936 KB |
testcase_22 | AC | 856 ms
152,320 KB |
testcase_23 | AC | 681 ms
171,440 KB |
testcase_24 | AC | 690 ms
169,028 KB |
testcase_25 | AC | 620 ms
172,492 KB |
testcase_26 | AC | 621 ms
174,408 KB |
testcase_27 | AC | 599 ms
174,488 KB |
testcase_28 | AC | 915 ms
150,516 KB |
testcase_29 | AC | 894 ms
150,400 KB |
testcase_30 | AC | 944 ms
149,760 KB |
ソースコード
class Tree(): def __init__(self, n, decrement=1): self.n = n self.edges = [[] for _ in range(n)] self._edge_label = [[] for _ in range(n)] self.root = None self.size = [1]*n # number of nodes in subtree self.decrement = decrement def add_edge(self, u, v, i): u, v = u-self.decrement, v-self.decrement self.edges[u].append(v) self.edges[v].append(u) self._edge_label[u].append((v,i)) self._edge_label[v].append((u,i)) def add_edges(self, edges): for i, p in enumerate(edges): u, v = p u, v = u-self.decrement, v-self.decrement self.edges[u].append(v) self.edges[v].append(u) self._edge_label[u].append((v, i)) self._edge_label[v].append((u, i)) def set_root(self, root): root -= self.decrement self.depth = [-1]*self.n self.root = root self.par = [-1]*self.n self.depth[root] = 0 self.edge_label = [-1]*self.n self.order = [root] next_set = [root] while next_set: p = next_set.pop() for q, i in self._edge_label[p]: if self.depth[q] != -1: continue self.par[q] = p self.depth[q] = self.depth[p]+1 self.edge_label[q]=i self.order.append(q) next_set.append(q) for p in self.order[::-1]: for q in self.edges[p]: if self.par[p] == q: continue self.size[p] += self.size[q] def diameter(self, path=False): # assert self.root is not None u = self.depth.index(max(self.depth)) dist = [-1]*self.n dist[u] = 0 prev = [-1]*self.n next_set = [u] while next_set: p = next_set.pop() for q in self.edges[p]: if dist[q] != -1: continue dist[q] = dist[p]+1 prev[q] = p next_set.append(q) d = max(dist) if path: v = w = dist.index(d) path = [v+1] while w != u: w = prev[w] path.append(w+self.decrement) return d, v+self.decrement, u+self.decrement, path else: return d def rerooting(self, op, merge, id): # assert self.root is not None dp1 = [id] * self.n dp2 = [id] * self.n for p in self.order[::-1]: t = id for q in self.edges[p]: if self.par[p] == q: continue dp2[q] = t t = merge(t, op(dp1[q], p, q)) t = id for q in self.edges[p][::-1]: if self.par[p] == q: continue dp2[q] = merge(t, dp2[q]) t = merge(t, op(dp1[q], p, q)) dp1[p] = t for q in self.order[1:]: pq = self.par[q] dp2[q] = op(merge(dp2[q], dp2[pq]), q, pq) dp1[q] = merge(dp1[q], dp2[q]) return dp1 def heavy_light_decomposition(self): """ return flat array of lists of heavy edges (1-indexed if decrement=True) """ # assert self.root is not None self.vid = [-1]*self.n self.hld = [-1]*self.n self.head = [-1]*self.n self.head[self.root] = self.root self.heavy_node = [-1]*self.n next_set = [self.root] for i in range(self.n): """ for tree graph, dfs ends in N times """ p = next_set.pop() self.vid[p] = i self.hld[i] = p+self.decrement maxs = 0 for q in self.edges[p]: """ encode direction of Heavy edge into heavy_node """ if self.par[p] == q: continue if maxs < self.size[q]: maxs = self.size[q] self.heavy_node[p] = q for q in self.edges[p]: """ determine "head" of heavy edge """ if self.par[p] == q or self.heavy_node[p] == q: continue self.head[q] = q next_set.append(q) if self.heavy_node[p] != -1: self.head[self.heavy_node[p]] = self.head[p] next_set.append(self.heavy_node[p]) return self.hld def lca(self, u, v): # assert self.head is not None u, v = u-self.decrement, v-self.decrement while True: if self.vid[u] > self.vid[v]: u, v = v, u if self.head[u] != self.head[v]: v = self.par[self.head[v]] else: return u + self.decrement def path(self, u, v): """ return the path array of the shortest path on u-v """ p = self.lca(u, v) u, v, p = u-self.decrement, v-self.decrement, p-self.decrement R = [] while u != p: yield u+self.decrement u = self.par[u] yield p+self.decrement while v != p: R.append(v) v = self.par[v] for v in reversed(R): yield v+self.decrement def distance(self, u, v): # assert self.head is not None p = self.lca(u, v) u, v, p = u-self.decrement, v-self.decrement, p-self.decrement return self.depth[u] + self.depth[v] - 2*self.depth[p] def find(self, u, v, x): return self.distance(u,x)+self.distance(x,v)==self.distance(u,v) def path_to_list(self, u, v, edge_query=False): """ transform a half-open interval into segments on the self.hld, which is the heavy edge list edge_query: map from edge (par,chi) to point (chi) (note: The root is never updated) """ # assert self.head is not None u, v = u-self.decrement, v-self.decrement while True: if self.vid[u] > self.vid[v]: u, v = v, u if self.head[u] != self.head[v]: yield self.vid[self.head[v]], self.vid[v] + 1 v = self.par[self.head[v]] else: yield self.vid[u] + edge_query, self.vid[v] + 1 return def ver_to_idx(self, u): """ return index on self.hld corresponding to vertex u """ return self.vid[u-self.decrement] def idx_to_ver(self, i): """ from index i on self.hld to vertex-index """ return self.hld[i] def idx_to_edge(self, i): """ from index i on self.hld to edge-index """ return self.edge_label[self.hld[i]-self.decrement] def subtree_query(self, u): u -= self.decrement return self.vid[u], self.vid[u] + self.size[u] def top_down(self,dp): for par in self.order: dp[par]=self.size[par]**2 for chi in self.edges[par]: if chi==self.par[par]: continue dp[par]-=self.size[chi]**2 return dp def top_down_edge_query(self,dp): def merge(dp_chi,dp_par): return dp_chi^dp_par for p in self.order[1+len(self.edges[self.root]):]: chi,par=self.edge_label[p],self.edge_label[self.par[p]] dp[chi]=merge(dp[chi], dp[par]) return dp def bottom_up(self,dp): for par in self.order[::-1]: tmp=0 for chi in self.edges[par]: if self.par[par] == chi: continue tmp+=dp[chi] dp[par]=(self.size[par])*(self.size[par])-tmp print(dp) return dp def draw(self): import matplotlib.pyplot as plt import networkx as nx G = nx.Graph() for x in range(self.n): for y in self.edges[x]: G.add_edge(x + self.decrement, y + self.decrement) pos = nx.spring_layout(G) nx.draw_networkx(G, pos) plt.axis("off") plt.show() ######################################################################################################### def example(): global input example = iter( """ 3 1 2 1 3 """ .strip().split("\n")) input = lambda: next(example) ######################################################################################################### import sys input = sys.stdin.readline # example() N=int(input()) T = Tree(N,decrement=1) for i in range(N-1): x, y = map(int, input().split()) T.add_edge(x,y,i) T.set_root(1) dp=[0]*N T.top_down(dp) print(*dp,sep="\n")