結果

問題 No.2949 Product on Tree
ユーザー Carpenters-Cat
提出日時 2024-10-25 23:07:50
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 542 ms / 2,000 ms
コード長 769 bytes
コンパイル時間 5,853 ms
コンパイル使用メモリ 253,640 KB
最終ジャッジ日時 2025-02-24 23:52:13
ジャッジサーバーID
(参考情報)
judge1 / judge3
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 46
権限があれば一括ダウンロードができます

ソースコード

diff #

#include <bits/stdc++.h>
using namespace std;
#include <atcoder/all>
using namespace atcoder;
using mint = modint998244353;
int N;
using Vi = std::vector<int> ;
using VVi = vector<Vi>;
using Vm = vector<mint>;
Vm dp1, dp2;
VVi gr;
Vm A;
void dfs(int s = 0, int p = -1) {
	dp1[s] = 0;
	dp2[s] = A[s];
	mint sum = 1;
	for (int& v : gr[s]) {
		if (v == p) continue;
		dfs(v, s);
		dp1[s] += dp1[v];
		dp2[s] += dp2[v] * A[s];
		dp1[s] += sum * dp2[v] * A[s];
		sum += dp2[v];
	}
}
int main () {
	cin >> N;
	gr.resize(N);
	A.resize(N);
	dp1.resize(N); dp2.resize(N);
	for (auto& a : A) {
		int x;
		cin >> x;
		a = x;
	}
	for (int i = 1; i < N; i ++) {
		int a, b;
		cin >> a >> b;
		gr[--a].push_back(--b);
		gr[b].push_back(a);
	}
	dfs();
	cout << dp1[0].val() << endl;
}
0