poj 1741 tree（点分治）-优快云博客

本文链接：https://blog.youkuaiyun.com/Kuguotao/article/details/89452513

给定一棵有n个节点的树，每条边具有不超过1000的长度。有效对(u, v)是指节点u和v之间的距离不超过给定的整数k。任务是编写程序计算满足条件的节点对数量。输入包含多组测试用例，每组用例包括n（n≤10000）、k以及n-1条边及其长度。最后一组测试用例由两个零结束。输出要求对每组用例在一行中输出答案。" 118812254,5891101,SQL分数排名查询,"['数据库', 'SQL查询', '数据处理']

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Description

Give a tree with n vertices,each edge has a length(positive integer less than 1001).
Define dist(u,v)=The min distance between node u and v.
Give an integer k,for every pair (u,v) of vertices is called valid if and only if dist(u,v) not exceed k.
Write a program that will count how many pairs which are valid for a given tree.

Input

The input contains several test cases. The first line of each test case contains two integers n, k. (n<=10000) The following n-1 lines each contains three integers u,v,l, which means there is an edge between node u and v of length l.
The last test case is followed by two zeros.

Output

For each test case output the answer on a single line.

Sample Input

Sample Output

/*

*/
#include <cstdio>
#include <queue>
#include <cstring>
#include <iostream>
#include <algorithm>
#include <math.h>
#include <map>
#include <set>
#include <bitset>
#include <stack>
#define ull unsigned long long
using namespace std;
typedef long long ll;
const ll mod=1000000007;
const int N=10006;
const double pi=acos(-1);
const int inf=0x3f3f3f3f;
struct node
{
    int to,c,next;
} g[N*2];
int head[N];
int son[N],f[N];//f记录以v为根的最大子树的大小
bool vis[N];
int d[N],deep[N];
int n,m,sum,root,ans,k;//k从一开始

void add_edge(int from,int to,int cost)
{

    g[k].to = to;
    g[k].c = cost;
    g[k].next = head[from];
//    printf("%d****\n",k);
    head[from] = k++;
}

void getroot(int v,int fa)//找重心
{
    son[v] = 1;
    f[v] = 0;//f记录以v为根的最大子树的大小
    for(int i = head[v]; i; i=g[i].next)
        if(g[i].to != fa && !vis[g[i].to])
        {
            getroot(g[i].to,v);//递归更新
            son[v] += son[g[i].to];
            f[v] = max(f[v],son[g[i].to]);//比较每个子树
        }
    f[v] = max(f[v],sum - son[v]);//别忘了以v父节点为根的子树

    if(f[v] < f[root])
        root = v;//更新当前根
}

void getdeep(int v,int fa)//获取子树所有节点与根的距离
{
    deep[++deep[0]] = d[v];
    for(int i = head[v]; i; i=g[i].next)
        if(g[i].to != fa && !vis[g[i].to])
        {
            d[g[i].to] = d[v] + g[i].c;
            getdeep(g[i].to,v);
        }
}

int cal(int v,int cost)//计算当前以重心v的子树下，所有情况的答案
{
    d[v] = cost;
    deep[0] = 0;
    getdeep(v,0);
    sort(deep+1,deep+deep[0]+1);
    int l = 1,r = deep[0],sum = 0;
    while(l < r)
    {
        if(deep[l]+deep[r] <= m)
        {
            sum += r-l;
            l++;
        }
        else
            r--;
    }
    return sum;
}

void solve(int v)//以v为重心进行计算
{
    ans += cal(v,0);
    vis[v] = 1;//标记
    for(int i = head[v]; i; i=g[i].next)
        if(!vis[g[i].to])
        {
            ans -= cal(g[i].to,g[i].c);
            sum = son[g[i].to];
            root = 0;
            getroot(g[i].to,0);
            solve(root);//递归处理下一个连通块
        }
}

int main()
{
    int u,v,w;
    while(~scanf("%d%d",&n,&m))
    {
        if(n==0&&m==0)
            return 0;
//        ans = root = k = 0;
        ans=root=0;k=1;
        memset(vis,0,sizeof(vis));
        memset(head,0,sizeof(head));
        for(int i = 1; i < n; i++)
        {
            scanf("%d%d%d",&u,&v,&w);
            add_edge(u,v,w);
            add_edge(v,u,w);
        }
        f[0] = inf;
        sum = n;
        getroot(1,0);
        solve(root);
        printf("%d\n",ans);
    }
    return 0;
}