因为求的路径都是从在包含根的一条链上。所以可以考虑直接dfs。记录栈里元素的前缀和。
比如搜到u了,在set里面查询有没有d[u]-k,有的话就加一。然后继续往下dfs的时候把d[u]插入set。
别忘了最开始插个0。
10w个点递归会爆?最好写个手工栈?哈哈我会写了!
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | #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; } |