【CF580C】Kefa and Park

题目大意：给定一棵 N 个节点的有根树（其中根节点始终为 1 号节点），点有点权，点权只有 1 和 0 两种，求从根节点到叶子节点的路径中，有多少条路径满足：路径上最大连续点权为 1 的节点个数不超过 M。

题解：一开始想到了树形dp，即：\(dp[i][0/1]\) 表示从根节点到该点的路径中不选/选该点的最大连续 1 的个数，则状态的转移涉及到该节点的 \(val\) 值和其父节点的 \(val\) 值。不过，其实没有必要记录下每一个节点的状态，原因是：如果到当前节点的连续 1 的个数已经超过了限制，那么后面无论怎样，这条路径都对答案没有贡献。因此，可以直接记录下到当前节点连续的长度，超过限制直接剪枝，由剪枝操作还可以知道，若当前节点的权值是 0，无论前面的值是多少，都对后面的值没有任何影响。

另外：建树的时候应该添加双向边，因为题目数据中给的关系并不是严格的父子关系。

树形 dp 代码如下

#include <bits/stdc++.h>
using namespace std;
const int maxn=1e5+10;

inline int read(){
	int x=0,f=1;char ch;
	do{ch=getchar();if(ch=='-')f=-1;}while(!isdigit(ch));
	do{x=x*10+ch-'0';ch=getchar();}while(isdigit(ch));
	return f*x;
}

struct node{
	int nxt,to;
}e[maxn<<1];
int tot=1,head[maxn];
inline void add_edge(int from,int to){
	e[++tot]=node{head[from],to},head[from]=tot;
}
int n,m,ans,val[maxn],dp[maxn][2];

void read_and_parse(){
	n=read(),m=read();
	for(int i=1;i<=n;i++)val[i]=read();
	for(int i=1;i<n;i++){
		int from=read(),to=read();
		add_edge(from,to),add_edge(to,from);
	}
}

void dfs(int u,int fa){
	if(!val[u])dp[u][0]=max(dp[fa][0],dp[fa][1]);
	else if(!val[fa])dp[u][1]=1,dp[u][0]=dp[fa][0];
	else dp[u][1]=dp[fa][1]+1,dp[u][0]=max(dp[fa][1],dp[fa][0]);
	int ok=1;
	for(int i=head[u];i;i=e[i].nxt){
		int v=e[i].to;if(v==fa)continue;
		ok=0,dfs(v,u);
	}
	if(ok&&max(dp[u][0],dp[u][1])<=m)++ans;
}

void solve(){
	dfs(1,0);
	printf("%d\n",ans);
}

int main(){
	read_and_parse();
	solve();
	return 0;
}

优化后代码如下

#include <bits/stdc++.h>
using namespace std;
const int maxn=1e5+10;

inline int read(){
	int x=0,f=1;char ch;
	do{ch=getchar();if(ch=='-')f=-1;}while(!isdigit(ch));
	do{x=x*10+ch-'0';ch=getchar();}while(isdigit(ch));
	return f*x;
}

struct node{
	int nxt,to;
}e[maxn<<1];
int tot=1,head[maxn];
inline void add_edge(int from,int to){
	e[++tot]=node{head[from],to},head[from]=tot;
}
int n,m,ans,val[maxn];

void read_and_parse(){
	n=read(),m=read();
	for(int i=1;i<=n;i++)val[i]=read();
	for(int i=1;i<n;i++){
		int from=read(),to=read();
		add_edge(from,to),add_edge(to,from);
	}
}

void dfs(int u,int fa,int cnt){
	if(cnt>m)return;
	bool flag=1;
	for(int i=head[u];i;i=e[i].nxt){
		int v=e[i].to;if(v==fa)continue;
		flag=0;
		if(val[v])dfs(v,u,cnt+1);
		else dfs(v,u,0);
	}
	ans+=flag;
}

void solve(){
	dfs(1,0,val[1]);
	printf("%d\n",ans);
}

int main(){
	read_and_parse();
	solve();
	return 0;
}