結果
問題 | No.2949 Product on Tree |
ユーザー |
![]() |
提出日時 | 2024-10-25 22:09:01 |
言語 | C++23 (gcc 13.3.0 + boost 1.87.0) |
結果 |
AC
|
実行時間 | 539 ms / 2,000 ms |
コード長 | 1,709 bytes |
コンパイル時間 | 3,596 ms |
コンパイル使用メモリ | 257,888 KB |
実行使用メモリ | 60,492 KB |
最終ジャッジ日時 | 2024-10-25 22:10:09 |
合計ジャッジ時間 | 26,771 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge1 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 3 |
other | AC * 46 |
ソースコード
#include <bits/stdc++.h> using namespace std; using ll = long long; const int INF = 1e9 + 10; const ll INFL = 4e18; #include <atcoder/modint> using mint = atcoder::modint998244353; int main() { int N; cin >> N; vector<int> A(N); for (int i = 0; i < N; i++) cin >> A[i]; vector<vector<int>> G(N); for (int i = 0; i < N - 1; i++) { int u, v; cin >> u >> v; u--; v--; G[u].push_back(v); G[v].push_back(u); } vector<vector<mint>> dp(N); auto dfs1 = [&](auto dfs1, int now, int pre) -> mint { dp[now] = vector<mint>(ssize(G[now])); mint ret = 0; for (int i = 0; i < ssize(G[now]); i++) { int nxt = G[now][i]; if (nxt != pre) { dp[now][i] = dfs1(dfs1, nxt, now) * A[now]; ret += dp[now][i]; } } return ret + A[now]; }; mint ans = 0; auto dfs2 = [&](auto dfs2, int now, int pre, mint prp) -> void { int deg = ssize(G[now]); for (int i = 0; i < deg; i++) { if (G[now][i] == pre) dp[now][i] = prp * A[now]; } vector<mint> dpl(deg + 1, 0), dpr(deg + 1, 0); for (int i = 0; i < deg; i++) dpl[i + 1] = dpl[i] + dp[now][i]; for (int i = deg - 1; i >= 0; i--) dpr[i] = dpr[i + 1] + dp[now][i]; ans += dpl.back(); for (int i = 0; i < deg; i++) { int nxt = G[now][i]; if (nxt != pre) { mint nprp = dpl[i] + dpr[i + 1] + A[now]; dfs2(dfs2, nxt, now, nprp); } } }; dfs1(dfs1, 0, -1); dfs2(dfs2, 0, -1, 0); ans /= 2; cout << ans.val() << endl; }