\(\\\)
\(Description\)
一棵以\(1\)号节点为根的\(N\)个节点的树,每个节点有正整数点权\(A_i\),问有多少条路径的节点点权之和达到\(S\)。
此处定义的路径中,节点的深度必须是升序的,不能出现由两个点的\(Lca\)连接的两段路径,路径起点是任意的。
- \(N\in [1,10^5]\),\(A_i,S\in [0,10^3]\)
\(\\\)
\(Solution\)
树上前缀和的做法显然,记录每个点的树上前缀和,倍增处理\(2^k\)的祖先。枚举每一个点作为路径的终点,二分起点祖先层数,二进制拆分跳到那里然后前缀和相减验证,总复杂度\(\text O(Nlog^2N)\)。
考虑将两个部分合在一起。二分祖先层数是因为不知道这一部分的链上点权和,考虑树上倍增。预处理\(f[i][k]\)表示节点\(i\)的\(2^k\)层祖先是谁,处理\(g[i][k]\)表示从节点\(i\)向上长度为\(2^k\)的链上点权和\((\)不包括\(i\),即从\(i\)的父亲到\(i\)的\(2^k\)层祖先的链上所有点权和\()\)。
在每一个点尝试向上跳,维护当前路径上权值\(now\)和当前节点位置\(x\),若存在\(f[x][k]\)且\(now+g[x][k]\le S\)就往上跳,判断最后\(now\)是否与\(S\)相等即可,总复杂度\(\text O(NlogN)\)。
\(\\\)
\(Code\)
#include<cmath>
#include<queue>
#include<cstdio>
#include<cctype>
#include<cstdlib>
#include<cstring>
#include<iostream>
#include<algorithm>
#define N 100010
#define R register
#define gc getchar
using namespace std;
inline int rd(){
int x=0; bool f=0; char c=gc();
while(!isdigit(c)){if(c=='-')f=1;c=gc();}
while(isdigit(c)){x=(x<<1)+(x<<3)+(c^48);c=gc();}
return f?-x:x;
}
int n,m,t,tot,hd[N],val[N];
int ans,d[N],f[N][20],g[N][20];
struct edge{int to,nxt;}e[N<<1];
inline void add(int u,int v){
e[++tot].to=v; e[tot].nxt=hd[u]; hd[u]=tot;
}
queue<int> q;
inline void bfs(){
q.push(1); d[1]=1;
while(!q.empty()){
int u=q.front(); q.pop();
for(R int i=hd[u],v;i;i=e[i].nxt)
if(!d[v=e[i].to]){
d[v]=d[u]+1;
f[v][0]=u; g[v][0]=val[u];
for(R int i=1;i<=t;++i){
f[v][i]=f[f[v][i-1]][i-1];
g[v][i]=g[v][i-1]+g[f[v][i-1]][i-1];
}
q.push(v);
}
}
}
inline bool valid(int x){
int now=val[x];
for(R int i=t;~i;--i)
if(f[x][i]&&now+g[x][i]<=m) now+=g[x][i],x=f[x][i];
return now==m;
}
int main(){
t=log2(n=rd()); m=rd();
for(R int i=1;i<=n;++i) val[i]=rd();
for(R int i=1,u,v;i<n;++i){
u=rd(); v=rd();
add(u,v); add(v,u);
}
bfs();
for(R int i=1;i<=n;++i) ans+=valid(i);
printf("%d\n",ans);
return 0;
}