結果
問題 | No.2949 Product on Tree |
ユーザー |
![]() |
提出日時 | 2024-10-25 22:33:11 |
言語 | PyPy3 (7.3.15) |
結果 |
AC
|
実行時間 | 575 ms / 2,000 ms |
コード長 | 1,108 bytes |
コンパイル時間 | 331 ms |
コンパイル使用メモリ | 82,488 KB |
実行使用メモリ | 132,380 KB |
最終ジャッジ日時 | 2024-10-25 22:33:37 |
合計ジャッジ時間 | 25,383 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge1 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 3 |
other | AC * 46 |
ソースコード
import sys input = sys.stdin.readline N = int(input()) A = list(map(int, input().split())) G = [[] for i in range(N)] for i in range(N - 1): u, v = map(int, input().split()) u, v = u - 1, v - 1 G[u].append(v) G[v].append(u) mod = 998244353 def non_rec_dfs(s, G): N = len(G) stack = [s] par = [-1] * N sz = [1] * N dp = [0] * N ans = 0 while stack: u = stack.pop() if u >= 0: stack.append(~u) for v in G[u]: if v == par[u]: continue par[v] = u stack.append(v) else: u = ~u s1, s2 = 0, 0 for v in G[u]: if v == par[u]: continue sz[u] += sz[v] dp[u] += dp[v] * A[u] dp[u] %= mod s1 += dp[v] s2 += dp[v] ** 2 dp[u] += A[u] dp[u] %= mod ans += (s1 ** 2 - s2) // 2 * A[u] ans %= mod return (ans + sum(dp) - sum(A)) % mod print(non_rec_dfs(0, G))