POJ - 1741 - Tree - 点分治 模板

#include <algorithm>
#include  <iterator>
#include  <iostream>
#include   <cstring>
#include   <cstdlib>
#include   <iomanip>
#include    <bitset>
#include    <cctype>
#include    <cstdio>
#include    <string>
#include    <vector>
#include     <cmath>
#include     <queue>
#include      <list>
#include       <map>
#include       <set>
using namespace std;
//#pragma GCC optimize(3)
//#pragma comment(linker, "/STACK:102400000,102400000")  //c++
#define lson (l , mid , rt << 1)
#define rson (mid + 1 , r , rt << 1 | 1)
#define debug(x) cerr << #x << " = " << x << "\n";
#define pb push_back
#define pq priority_queue



typedef long long ll;
typedef unsigned long long ull;

typedef pair<ll ,ll > pll;
typedef pair<int ,int > pii;
typedef pair<int,pii> p3;

//priority_queue<int> q;//这是一个大根堆q
//priority_queue<int,vector<int>,greater<int> >q;//这是一个小根堆q
#define fi first
#define se second
//#define endl ‘\n‘

#define OKC ios::sync_with_stdio(false);cin.tie(0)
#define FT(A,B,C) for(int A=B;A <= C;++A)  //用来压行
#define REP(i , j , k)  for(int i = j ; i <  k ; ++i)
//priority_queue<int ,vector<int>, greater<int> >que;

const ll mos = 0x7FFFFFFFLL;  //2147483647
const ll nmos = 0x80000000LL;  //-2147483648
const int inf = 0x3f3f3f3f;
const ll inff = 0x3f3f3f3f3f3f3f3fLL; //18
const int mod = 998244353;

const double PI=acos(-1.0);

// #define _DEBUG;         //*//
#ifdef _DEBUG
freopen("input", "r", stdin);
// freopen("output.txt", "w", stdout);
#endif
/*-----------------------showtime----------------------*/
            const int maxn = 1e5+9;
            int root = 0,S,mx;
            int n,k;
            int sz[maxn],f[maxn],dis[maxn],cnt;
            bool used[maxn];
            struct node
            {
                int to,w,nx;
            }e[maxn];
            int h[maxn],tot = 0;
            void add(int u,int v,int w){
                e[tot].to = v;
                e[tot].w = w;
                e[tot].nx = h[u];
                h[u] = tot++;
            }
            void getRoot(int u, int fa){
                sz[u] = 1,f[u] = 1;
                for(int i = h[u] ; ~i; i= e[i].nx){
                    int v = e[i].to;
                    if(used[v] || fa == v)continue;
                    getRoot(v,u);
                    sz[u] += sz[v];
                    f[u] = max(f[u] , sz[v]);
                }
                f[u] = max(f[u],S - sz[u]);
                if(f[u] < mx){root = u;mx = f[u];}
            }

            void getDis(int u,int fa,int D){
                for(int i=h[u] ; ~i; i=e[i].nx){
                    int v = e[i].to;
                    if(used[v]||v == fa)continue;
                    dis[++cnt] = D + e[i].w;
                    getDis(v,u,dis[cnt]);
                }
            }

            int getAns(int x,int D){
                dis[cnt = 1] = D;
                getDis(x,0,D);
                sort(dis+1,dis+1+cnt);
                int le = 1,ri =cnt,ans = 0;
                while(le <= ri){
                    if(dis[le] + dis[ri] <= k)ans += ri - le,le++;
                    else ri--;
                }
                return ans;
            }

            int Divide(int x){
                used[x] = true;
                ll ans = getAns(x,0);
                for(int i=h[x]; ~i; i= e[i].nx){
                    int v = e[i].to;
                    if(used[v])continue;
                    ans -= getAns(v,e[i].w);
                    mx = inf,S = sz[v];
                    getRoot(v,x);ans += Divide(root);
                }
                return ans;
            }
int main(){
            
            while(~scanf("%d%d", &n, &k) && n+k)
            {
                memset(h,-1,sizeof(h));
                memset(used,false,sizeof(used));
                tot = 0;
                for(int i=1; i<n; i++){
                    int u,v,c;
                    scanf("%d%d%d", &u, &v,&c);
                    add(u,v,c);
                    add(v,u,c);
                }
                S = n;mx = inf;
                getRoot(1,-1);
                printf("%d\n",Divide(root));
            }
            return 0;
}