CF1929D Sasha and a Walk in the City 题解

CF1929D Sasha and a Walk in the City

简单树形动态规划。

我们把选取到的点称为黑点，由题意得，一个合法的点集能使树中任意一条简单路径上的黑点数量不超过两个。也就是说，如果黑点数量多于 \(2\)，对于任意两个黑点，它们如果在同一个节点的子树内，必然是兄弟关系。否则，一旦存在祖先关系，由于黑点数量多于 \(2\)，必然有一个黑点可以与这两个点组成一条黑点数量超过两个的简单路径。

设状态 \(dp_{i,j}\) 表示以 \(i\) 为根的子树内，从根到叶子节点最多经过 \(j\) 个黑点。显然，只有 \(j\) 等于 \(0,1\) 或 \(2\) 时，状态是合法的。

考虑如何转移，对于 \(dp_{i,0}\)，显然等于 \(1\)。

对于 \(dp_{i,1}\)，由于根到叶子节点最多经过黑点数最大值为 \(1\)，最多就是两条到根的路径合并，经过 \(1+1=2\) 个黑点，所以每一个儿子内可以任意选择。另外，每个子树内还可以从没有黑点涂子树的根上的黑点变成有 \(1\) 个黑点，所以还需要加入没有黑点的方法数。根据乘法原理，可以推出如下转移式：

\[dp_{i,1}=\prod_{j\in son(i)}(dp_{j,0}+dp_{j,1})=\prod_{j\in son(i)}(1+dp_{j,1})
\]

对于 \(dp_{i,2}\)，由于根到叶子节点最多经过黑点数最大值为 \(2\)，所以只能有一个子树内的 \(dp\) 值可以转移来，否则必然可以构造一条黑点数量超过两个的简单路径。另外，每个子树内还可以从 \(1\) 个黑点涂子树的根上的黑点变成有 \(2\) 个黑点，所以还需要加入 \(1\) 个黑点的方法数。根据加法原理，可以推出如下转移式：

\[dp_{i,2}=\sum_{j\in son(i)}(dp_{j,1}+dp_{j,2})
\]

以 \(1\) 为整棵树的根，最后的答案为 \(dp_{1,0}+dp_{1,1}+dp_{1,2}\)，也就是 \(1+dp_{1,1}+dp_{1,2}\)。

#include <bits/stdc++.h>
using namespace std;
struct edge
{
	long long v,nxt;
}e[800000];
long long t,u,v,n,mod=998244353,h[800000],f[800000][3],cnt=0;
void add_edge(long long u,long long v)
{
	e[++cnt].nxt=h[u];
	e[cnt].v=v;
	h[u]=cnt;
}

void init()
{
	for(int i=1;i<=n;i++)h[i]=0,f[i][1]=1,f[i][2]=0;
	cnt=0;
}

void dfs(long long x,long long fa)
{
	for(int i=h[x];i;i=e[i].nxt)
	    if(e[i].v!=fa)
	       {
		   dfs(e[i].v,x);
		   f[x][1]=(f[x][1]*(f[e[i].v][1]+1)%mod)%mod;
		   f[x][2]=(f[x][2]+f[e[i].v][1]+f[e[i].v][2])%mod;
	       }
}

int main()
{
	scanf("%lld",&t);
	while(t--)
	   {
	   	scanf("%lld",&n);
	   	init();
	   	for(int i=1;i<=n-1;i++)
	   	    {
	   	    scanf("%lld%lld",&u,&v);
	   	    add_edge(u,v),add_edge(v,u);
			}
		dfs(1,0);
		printf("%lld\n",(f[1][1]+f[1][2]+1)%mod);
	   }
	return 0;
}