题目大意：给定一棵 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;
}