牛客训练赛55 E 树

很妙的一个树形DP问题，简单考虑了一下就过了

https://ac.nowcoder.com/acm/contest/2927/E

主要就是推公式（公式有点长呀）

大概就是这样，其实挺简单的。

#include<iostream>
#include<cstring>
#include<algorithm>
#include<vector>
using namespace std;
typedef long long  ll;
const int maxn = 1e6 + 7;
const ll mod = 998244353;
int head[maxn];
int cnt = 1;
ll sum[maxn];
ll son[maxn];
ll dp[maxn];

struct Node {
	int p;
	int next;
}G[2*maxn];
int n;
void add(int be, int en) {
	G[cnt].p = en;   G[cnt].next = head[be];  head[be] = cnt;
	cnt++;
}
int dfs(int x,int fa) {
	son[x] = 1;
	for (int i = head[x]; i; i = G[i].next) {
		int p = G[i].p;
		if (p == fa) continue;
		dfs(p, x);
		son[x] += son[p];
		sum[x] = ( sum[x] + son[p] + sum[p]) % mod;
		dp[x] = (dp[x] + dp[p] + 2 * sum[p] + son[p]) % mod;
	}
	return 0;
}

int dfs2(int x, int fa) {
	for (int i = head[x]; i; i = G[i].next) {
		int p = G[i].p;
		if (p == fa) continue;

		ll ans = ((dp[x] - (dp[p] + 2 * sum[p] + son[p])) % mod + mod) % mod;//dp
		//sum
		ll cns = ((sum[x] - (sum[p] + son[p])) % mod + mod) % mod;

		dp[p] = ((dp[p] + ans + 2 * cns + n - son[p]) % mod + mod) % mod;
		sum[p] = ((sum[p] + cns + n - son[p]) % mod + mod) % mod;

		dfs2(p, x);
	}
	return 0;
}
int main() {
	int be, en;

	scanf("%d", &n);
	for (int i = 1; i < n; i++) {
		scanf("%d %d", &be, &en);
		add(be, en);
		add(en, be);
	}
	dfs(1, -1);
	dfs2(1, -1);
	ll a = 0;

	for (int i = 1; i <= n; i++) {
		a = (a + dp[i]) % mod;
	}
	printf("%lld\n", a);
	return 0;
}