BZOJ2783: [JLOI2012]树

  1 #include<cstdio>
  2 
  3 #include<cstdlib>
  4 
  5 #include<cmath>
  6 
  7 #include<cstring>
  8 
  9 #include<algorithm>
 10 
 11 #include<iostream>
 12 
 13 #include<vector>
 14 
 15 #include<map>
 16 
 17 #include<set>
 18 
 19 #include<queue>
 20 
 21 #include<string>
 22 
 23 #define inf 1000000000
 24 
 25 #define maxn 100000+5
 26 
 27 #define maxm 2000000+5
 28 
 29 #define eps 1e-10
 30 
 31 #define ll long long
 32 
 33 #define pa pair<int,int>
 34 
 35 #define for0(i,n) for(int i=0;i<=(n);i++)
 36 
 37 #define for1(i,n) for(int i=1;i<=(n);i++)
 38 
 39 #define for2(i,x,y) for(int i=(x);i<=(y);i++)
 40 
 41 #define for3(i,x,y) for(int i=(x);i>=(y);i--)
 42 
 43 #define for4(i,x) for(int i=head[x],y=e[i].go;i;i=e[i].next,y=e[i].go)
 44 
 45 #define for5(n,m) for(int i=1;i<=n;i++)for(int j=1;j<=m;j++)
 46 
 47 #define mod 1000000007
 48 
 49 #define mid ((l+r)>>1)
 50 
 51 #define lch k<<1,l,mid
 52 
 53 #define rch k<<1|1,mid+1,r
 54 
 55 using namespace std;
 56 
 57 inline int read()
 58 
 59 {
 60 
 61     int x=0,f=1;char ch=getchar();
 62 
 63     while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
 64 
 65     while(ch>='0'&&ch<='9'){x=10*x+ch-'0';ch=getchar();}
 66 
 67     return x*f;
 68 
 69 }
 70 int n,m,ti,cnt,tot,head[maxn],a[maxn],id[maxn],pos[maxn][2],sum,fa[maxn],rt[maxn];
 71 int s[maxm],ls[maxm],rs[maxm];
 72 struct edge{int go,next;}e[maxn];
 73 map<int,int>mp;
 74 inline void add(int x,int y)
 75 {
 76     e[++tot]=(edge){y,head[x]};head[x]=tot;
 77 }
 78 inline void dfs(int x)
 79 {
 80     pos[x][0]=++ti;id[ti]=x;
 81     for4(i,x)if(!pos[y][0])
 82     {
 83         fa[y]=x;a[y]+=a[x];
 84         dfs(y);
 85     }
 86     pos[x][1]=ti;
 87 }
 88 inline void update(int l,int r,int x,int &y,int z)
 89 {
 90     y=++sum;
 91     s[y]=s[x]+1;
 92     if(l==r)return;
 93     ls[y]=ls[x];rs[y]=rs[x];
 94     if(z<=mid)update(l,mid,ls[x],ls[y],z);else update(mid+1,r,rs[x],rs[y],z);
 95 }
 96 inline int query(int l,int r,int x,int y,int z)
 97 {
 98     if(l==r)return s[y]-s[x];
 99     return z<=mid?query(l,mid,ls[x],ls[y],z):query(mid+1,r,rs[x],rs[y],z);
100 }
101 
102 int main()
103 
104 {
105 
106     freopen("input.txt","r",stdin);
107 
108     freopen("output.txt","w",stdout);
109 
110     n=read();m=read();
111     for1(i,n)a[i]=read();
112     for1(i,n-1){int x=read(),y=read();add(x,y);}
113     dfs(1);
114     for1(i,n)if(!mp[a[i]])mp[a[i]]=++cnt;
115     for1(i,n)update(1,cnt,rt[i-1],rt[i],mp[a[id[i]]]);
116     ll ans=0;
117     for1(i,n)
118     {
119         int t=mp[a[fa[i]]+m];
120         if(t)ans+=query(1,cnt,rt[pos[i][0]-1],rt[pos[i][1]],t);
121     }
122     cout<<ans<<endl;
123 
124     return 0;
125 
126 }