P4178 Tree

题意：

给定一棵 n 个节点的树，每条边有边权，求出树上两点距离小于等于 k 的点对数量。

题解：

点分治的模板题是求等于K的路径条数
本题是求小于等于K的路径条数，我们只需要改变统计答案即可
原本统计答案是对一个路劲长度len，判断K-len在之前的子树中出现多少次，用f数组来记录，直接查看f[k-len]等于多少即可
现在统计答案是对一个路径长度len，判断小于等于K-len的数在之前的子树中出现多少次，统计区间数量，我们可以用树状数组来实现，对于每个len我们插入数组f[]其中，然后求1~K-len的区间查询

代码：

详细看代码

#include<bits/stdc++.h>
#define MAXN 40005
#define MAXK 20005
using namespace std;int N,K;struct edge{int v,w;edge(int v=0, int w=0):v(v), w(w){}
};vector<edge> adj[MAXN];//
int fw[MAXK];
int lbt(int x){return x & (-x);
}
int getsum(int x){int ans = 0;for(;x>0;x-=lbt(x)){ans += fw[x];}return ans;
}void change(int x, int dv){for(;x<=K;x+=lbt(x)){fw[x] += dv;}
}//
int sz[MAXN];
bool vis[MAXN];
int rt;
void dfs_rt(int u, int fa, int tot){//O(tot)求根 sz[u] = 1;int v, n = 0;for(int k=0;k<adj[u].size();k++){v = adj[u][k].v;if(v==fa || vis[v]) continue;dfs_rt(v, u, tot);sz[u] += sz[v];n = max(n, sz[v]);}n = max(n, tot-sz[u]);if(n*2 <= tot) rt = u;
}void dfs_sz(int u, int fa){//O(tot)求子树sz[u] = 1;int v, n = 0;for(int k=0;k<adj[u].size();k++){v = adj[u][k].v;if(v==fa || vis[v]) continue;dfs_sz(v, u);sz[u] += sz[v];}
}int d[MAXN], cnt = 0;
void dfs_dis(int u, int fa, int dis){//O(tot)d[++cnt] = dis;//记录距离 int v,w;for(int k=0;k<adj[u].size();k++){v = adj[u][k].v;w = adj[u][k].w;if(v==fa || vis[v]) continue;dfs_dis(v, u, dis + w);}
}void dfs_clear(int u, int fa, int dis){//O(tot)清零 if(dis) change(dis, -1);int v,w;for(int k=0;k<adj[u].size();k++){v = adj[u][k].v;w = adj[u][k].w;if(v==fa || vis[v]) continue;dfs_clear(v, u, dis + w);}
}int work(int u, int tot){   dfs_rt(u, 0, tot);u = rt;dfs_sz(u, 0);vis[u] = 1;int v,w;//solve/*求出每个点到根节点的距离，然后记录到数组d然后对于d[i],查询0~K-d[i]区间大小，用树状数组实现 如果d[i]<=K,那么d[i]本身也符合条件 将数组d插入到树状数组当以一个点为根的一轮结束时记得情况数组树状数组 */ int ans = 0;for(int k=0;k<adj[u].size();k++){v = adj[u][k].v;w = adj[u][k].w;if(vis[v]) continue;cnt = 0;dfs_dis(v, u, w);for(int i=1;i<=cnt;i++){if(d[i] <= K) ++ans;ans += getsum(max(0,K-d[i]));}for(int i=1;i<=cnt;i++){change(d[i], +1);}}dfs_clear(u,0,0);//手动清空 //for(int k=0;k<adj[u].size();k++){v = adj[u][k].v;//w = adj[u][k].w;if(vis[v]) continue;ans += work(v,sz[v]);}   return ans;
}int main(){cin>>N;int u,v,w;for(int i=1;i<N;i++){cin>>u>>v>>w;adj[u].push_back(edge(v,w));adj[v].push_back(edge(u,w));}cin>>K;cout<<work(1, N)<<endl;return 0;
}