結果
問題 | No.2949 Product on Tree |
ユーザー |
![]() |
提出日時 | 2024-10-25 22:30:26 |
言語 | PyPy3 (7.3.15) |
結果 |
AC
|
実行時間 | 1,986 ms / 2,000 ms |
コード長 | 1,907 bytes |
コンパイル時間 | 240 ms |
コンパイル使用メモリ | 82,400 KB |
実行使用メモリ | 279,788 KB |
最終ジャッジ日時 | 2024-10-25 22:31:39 |
合計ジャッジ時間 | 68,795 ms |
ジャッジサーバーID (参考情報) |
judge1 / judge3 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 3 |
other | AC * 46 |
ソースコード
from types import GeneratorType def bootstrap(f, stack=[]): def wrappedfunc(*args, **kwargs): if stack: return f(*args, **kwargs) to = f(*args, **kwargs) while True: if type(to) is GeneratorType: stack.append(to) to = next(to) else: stack.pop() if not stack: break to = stack[-1].send(to) return to return wrappedfunc N = int(input()) A = list(map(int, input().split())) G = [[] for _ in range(N)] for _ in range(N-1): u, v = map(int, input().split()) G[u-1].append(v-1) G[v-1].append(u-1) MOD = 998244353 @bootstrap def dfs(n, pre): ans = 0 for i, v in enumerate(G[n]): if v != pre: dp[n][i] = yield dfs(v, n) ans += dp[n][i] ans %= MOD yield ((ans*A[n])%MOD+A[n])%MOD dp = [[0]*len(G[i]) for i in range(N)] dfs(0, -1) def LRset(n): for i in range(1, len(G[n])): L[n][i] = (L[n][i-1]+dp[n][i-1])%MOD for i in reversed(range(len(G[n])-1)): R[n][i] = (R[n][i+1]+dp[n][i+1])%MOD L = [[0]*len(G[i]) for i in range(N)] R = [[0]*len(G[i]) for i in range(N)] LRset(0) @bootstrap def rerooting(n, pre, no): visited[n] = True preno = 0 for i in range(len(G[n])): if G[n][i] == pre: preno = i break dp[n][preno] = ((L[pre][no]+R[pre][no])*A[pre]%MOD+A[pre])%MOD LRset(n) for i in range(len(G[n])): if not visited[G[n][i]]: yield rerooting(G[n][i], n, i) yield visited = [False]*N visited[0] = True for i in range(len(G[0])): if not visited[G[0][i]]: rerooting(G[0][i], 0, i) ans = 0 for i, d in enumerate(dp): for n in d: ans += n*A[i] ans %= MOD def inverse(n, d): return n * pow(d, -1, MOD) % MOD print(inverse(ans, 2))