Codeforces Round 926 (Div. 2) D题 Sasha and a Walk in the City（树形dp）

题目链接

Codeforces Round 926 (Div. 2) D题 Sasha and a Walk in the City

思路

题意为计算不存在三点在一条线上的点集的数量。

考虑使用dp。

定义$dp[i][j]$表示以$i$为根的子树内，从根节点到叶子节点最多经过$j$个选取点。显然，$j$只有为$0$，$1$，$2$时是合法的。

当$j=0$时，$dp[i][0]$的答案显然为$1$。

当$j=1$时，最多的情况下便是两条到根的路径合并，最多经过$1+1=2$个选取点。所以每一个儿子内可以任意选择。因为可以将$i$节点这个根选上，所以还要加上没有选取点的方案数。则状态转移方程为：$dp[i][1]={\prod }_{j\in son(i)}^{}(dp[j][0]+dp[j][1])$。

当$j=2$时，由于根到叶子节点最多经过选取点数最大值为$2$，所以只能有一个子树内的dp值可以转移过来。因为根节点可以作为选取点，所以还要加上$1$个选取点的方案数。则状态转移方程为：$dp[i][2]={\sum }_{j\in son(i)}^{}(dp[j][1]+dp[j][2])$。

规定$1$为整棵树的根，则最终答案为$dp[i][0]+dp[i][1]+dp[i][2]$。

代码

#pragma GCC optimize("O2")
#pragma GCC optimize("O3")
#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N = 3e5 + 5, M = 1e6 + 5;
const int mod = 998244353;
const int inf = 0x3f3f3f3f3f3f3f3f;

int n;
int dp[N][3];
vector<int>mp[N];
void dfs(int u, int fu)
{
	dp[u][0] = 1;
	dp[u][1] = 1;
	for (int j : mp[u])
	{
		if (j == fu) continue;
		dfs(j, u);
		dp[u][1] = dp[u][1] * (dp[j][0] + dp[j][1]) % mod;
		dp[u][2] = dp[u][2] + (dp[j][1] + dp[j][2]) % mod;
	}
}
int MOD(int x)
{
	return (x % mod + mod) % mod;
}
void solve()
{
	cin >> n;
	for (int i = 1; i <= n; i++)
	{
		mp[i].clear();
		dp[i][0] = dp[i][1] = dp[i][2] = 0;
	}
	for (int i = 1, u, v; i < n; i++)
	{
		cin >> u >> v;
		mp[u].push_back(v);
		mp[v].push_back(u);
	}
	dfs(1, -1);
	int ans = 0;
	for (int i = 0; i <= 2; i++)
		ans = (ans + dp[1][i]) % mod;
	cout << MOD(ans) << endl;
}

signed main()
{
	ios::sync_with_stdio(false);
	cin.tie(0), cout.tie(0);
	int test = 1;
	cin >> test;
	for (int i = 1; i <= test; i++)
	{
		solve();
	}
	return 0;
}