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;
}