結果
問題 | No.2360 Path to Integer |
ユーザー |
![]() |
提出日時 | 2024-05-03 17:37:43 |
言語 | PyPy3 (7.3.15) |
結果 |
AC
|
実行時間 | 499 ms / 2,500 ms |
コード長 | 3,024 bytes |
コンパイル時間 | 285 ms |
コンパイル使用メモリ | 82,560 KB |
実行使用メモリ | 132,960 KB |
最終ジャッジ日時 | 2024-11-24 20:40:30 |
合計ジャッジ時間 | 5,584 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge3 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 2 |
other | AC * 15 |
ソースコード
import sysinput = sys.stdin.readlinemod = 998244353class ReRooting():def __init__(self, N):self.N = Nself.G = [[] for i in range(N)]self.mem = [[] for i in range(N)]self.AcL = [None] * Nself.AcR = [None] * Nself.par = [-1] * Nself.sz = [1] * Nself.unit = 0self.dp1 = [self.unit] * Nself.dp2 = [self.unit] * Ndef add_edge(self, i, j):self.G[i].append(j)self.G[j].append(i)def merge(self, a, b):return (a + b) % moddef calc_dp(self, v, s, sz):return (v * pow(10, len(str(s)), mod) + s * sz) % moddef calc_ans(self, v, s):return (v * pow(10, len(str(s)), mod) + s * N) % moddef bottom_up(self, s):stack = [s]while stack:u = stack.pop()if u >= 0:stack.append(~u)for v in self.G[u]:if v == self.par[u]:continueself.par[v] = ustack.append(v)else:u = ~uself.mem[u] = []for v in self.G[u]:if v == self.par[u]:continueself.sz[u] += self.sz[v]self.dp1[u] = self.merge(self.dp1[u], self.dp1[v])self.mem[u].append(v)self.dp1[u] = self.calc_dp(self.dp1[u], A[u], self.sz[u])Nt = len(self.mem[u])self.AcL[u] = [self.unit] * (Nt + 1)self.AcR[u] = [self.unit] * (Nt + 1)for i, v in enumerate(self.mem[u]):self.AcL[u][i + 1] = self.merge(self.AcL[u][i], self.dp1[v])for i, v in enumerate(self.mem[u][::-1]):self.AcR[u][i + 1] = self.merge(self.AcR[u][i], self.dp1[v])returndef top_down(self, s):stack = [s]while stack:u = stack.pop()for i, v in enumerate(self.mem[u]):self.dp2[v] = self.merge(self.merge(self.dp2[u], self.AcL[u][i]), self.AcR[u][len(self.mem[u]) - i - 1])self.dp2[v] = self.calc_dp(self.dp2[v], A[u], N - self.sz[v])stack.append(v)returndef solve(self, s):self.bottom_up(s)self.top_down(s)# print(self.dp1)# print(self.dp2)ans = [None] * self.Nfor i in range(self.N):ans[i] = self.dp2[i]for u in self.G[i]:if self.par[i] == u:continueans[i] = self.merge(ans[i], self.dp1[u])ans[i] = self.calc_ans(ans[i], A[i])return ansN = int(input())A = list(map(int, input().split()))G = ReRooting(N)for _ in range(N - 1):u, v = map(int, input().split())u, v = u - 1, v - 1G.add_edge(u, v)print(sum(G.solve(0))%mod)