因为求的路径都是从在包含根的一条链上。所以可以考虑直接dfs。记录栈里元素的前缀和。
比如搜到u了,在set里面查询有没有d[u]-k,有的话就加一。然后继续往下dfs的时候把d[u]插入set。
别忘了最开始插个0。
10w个点递归会爆?最好写个手工栈?哈哈我会写了!
#include<iostream> #include<cstdio> #include<cstring> #include<algorithm> #include<set> using namespace std; const int maxn=200005; int n,K,vis[maxn]={0},w[maxn],h[maxn],st[maxn],d[maxn],p[maxn],n1[maxn],tot=0; set<int>S; char buf[5000000],*pt=buf; inline int getint(){ int r=0;while(*pt<'0'||*pt>'9')pt++; while(*pt>='0'&&*pt<='9')r=r*10+*pt++-48;return r; }inline void ae(int a,int b){ p[tot]=b; n1[tot]=h[a]; h[a]=tot++; p[tot]=a; n1[tot]=h[b]; h[b]=tot++; }int main(){ fread(pt,1,5000000,stdin); n=getint(); K=getint(); memset(h,-1,sizeof(h)); S.clear(); S.insert(0); for(int i=1;i<=n;i++) w[i]=getint(); int u,v,An=0; for(int i=1;i<n;i++) u=getint(),v=getint(),ae(u,v); int top=0; st[++top]=1; S.insert(w[1]); d[1]=w[1]; vis[1]=1; while(top){ int o=st[top],flag=0; for(int i=h[o];~i;i=n1[i]) if(!vis[p[i]]){ st[++top]=p[i]; h[o]=n1[i]; d[p[i]]=d[o]+w[p[i]]; vis[p[i]]=flag=1; S.insert(d[p[i]]); if(S.find(d[p[i]]-K)!=S.end()) An++;break; }if(!flag){ top--; S.erase(d[o]);} }printf("%d\n",An); return 0; }