BZOJ2783: [JLOI2012]树【dfs+set】

因为求的路径都是从在包含根的一条链上。所以可以考虑直接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;
}