给你一棵树,以及这棵树上边的距离.问有多少对点它们两者间的距离小于等于K
点分治
我们如果开桶记录路径长度的数,那么需要维护一个单点加和前缀和,用树状数组维护就行了
不过还有种排序双指针的方法,复杂度一样,懒得写了qwq
复杂度(O(nlog^2n))
Code
#include <iostream>
#include <cstdio>
#include <algorithm>
#include <vector>
const int N = 4e4;
const int INF = 4e5;
using namespace std;
struct node
{
int to,cost;
};
int n,k,rt,maxp[N + 5],size[N + 5],su,vis[N + 5],ans,c[INF + 5],s[N + 5],num,cnt,now[N + 5];
vector <node> d[N + 5];
int lowbit(int x)
{
return x & (-x);
}
void add(int x,int s)
{
for (;x <= k;x += lowbit(x))
c[x] += s;
}
int query(int x)
{
int ans = 0;
for (;x;x -= lowbit(x))
ans += c[x];
return ans;
}
void get_rt(int u,int fa)
{
maxp[u] = 0;
size[u] = 1;
vector <node>::iterator it;
for (it = d[u].begin();it != d[u].end();it++)
{
int v = (*it).to;
if (v == fa || vis[v])
continue;
get_rt(v,u);
size[u] += size[v];
maxp[u] = max(maxp[u],size[v]);
}
maxp[u] = max(maxp[u],su - size[u]);
if (maxp[u] < maxp[rt])
rt = u;
}
void get_dis(int u,int fa,int s)
{
if (s > k)
return;
now[++cnt] = s;
vector <node>::iterator it;
for (it = d[u].begin();it != d[u].end();it++)
{
int v = (*it).to,w = (*it).cost;
if (vis[v] || v == fa)
continue;
get_dis(v,u,s + w);
}
}
void calc(int u)
{
num = 0;
vector <node>::iterator it;
for (it = d[u].begin();it != d[u].end();it++)
{
int v = (*it).to,w = (*it).cost;
if (vis[v])
continue;
cnt = 0;
get_dis(v,u,w);
for (int i = 1;i <= cnt;i++)
ans += query(k - now[i]) + 1;
for (int i = 1;i <= cnt;i++)
{
s[++num] = now[i];
add(now[i],1);
}
}
for (int i = 1;i <= num;i++)
add(s[i],-1);
}
void solve(int u)
{
vis[u] = 1;
calc(u);
vector <node>::iterator it;
for (it = d[u].begin();it != d[u].end();it++)
{
int v = (*it).to;
if (vis[v])
continue;
su = size[v];
rt = 0;
maxp[0] = INF;
get_rt(v,0);
solve(rt);
}
}
int main()
{
scanf("%d",&n);
int u,v,w;
for (int i = 1;i < n;i++)
{
scanf("%d%d%d",&u,&v,&w);
d[u].push_back((node){v,w});
d[v].push_back((node){u,w});
}
scanf("%d",&k);
maxp[0] = INF;
su = n;
get_rt(1,0);
solve(rt);
printf("%d
",ans);
return 0;
}