【POJ 1741】【BZOJ 1468】【BZOJ 3365】【点分治模板题】Tree

Description

给你一棵树，问有多少对点的距离小于等于K。

Solution

点分治裸题就没有什么好讲的了。。

Source

/************************************************
 * Au: Hany01
 * Date: Mar 2nd, 2018
 * Prob: POJ1741
 * Email: hany01@foxmail.com
************************************************/

#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<iostream>
#include<algorithm>
#include<cmath>

using namespace std;

typedef long long LL;
typedef pair<int, int> PII;
#define rep(i, j) for (register int i = 0, i##_end_ = (j); i < i##_end_; ++ i)
#define For(i, j, k) for (register int i = (j), i##_end_ = (k); i <= i##_end_; ++ i)
#define Fordown(i, j, k) for (register int i = (j), i##_end_ = (k); i >= i##_end_; -- i)
#define Set(a, b) memset(a, b, sizeof(a))
#define Cpy(a, b) memcpy(a, b, sizeof(a))
#define fir first
#define sec second
#define pb(a) push_back(a)
#define mp(a, b) make_pair(a, b)
#define ALL(a) (a).begin(), (a).end()
#define SZ(a) ((int)(a).size())
#define INF (0x3f3f3f3f)
#define INF1 (2139062143)
#define Mod (1000000007)
#define debug(...) fprintf(stderr, __VA_ARGS__)
#define y1 wozenmezhemecaia

template <typename T> inline bool chkmax(T &a, T b) { return a < b ? a = b, 1 : 0; }
template <typename T> inline bool chkmin(T &a, T b) { return b < a ? a = b, 1 : 0; }

inline int read()
{
    register int _, __; register char c_;
    for (_ = 0, __ = 1, c_ = getchar(); c_ < '0' || c_ > '9'; c_ = getchar()) if (c_ == '-') __ = -1;
    for ( ; c_ >= '0' && c_ <= '9'; c_ = getchar()) _ = (_ << 1) + (_ << 3) + (c_ ^ 48);
    return _ * __;
}

inline void File()
{
#ifdef hany01
    freopen("poj1741.in", "r", stdin);
    freopen("poj1741.out", "w", stdout);
#endif
}

const int maxn = 10005, maxm = 20005;

int n, k, v[maxm], w[maxm], nex[maxm], beg[maxn], e, sz[maxn], maxch[maxn], rt, Ans, dep[maxn], lst[maxn], cnt, nodesum, vis[maxn];

inline void add(int uu, int vv, int ww) { v[++ e] = vv, w[e] = ww, nex[e] = beg[uu], beg[uu] = e; }

void GetRoot(int u, int fa)
{
    sz[u] = 1, maxch[u] = 1;
    for (register int i = beg[u]; i; i = nex[i]) if (fa != v[i] && !vis[v[i]])
        GetRoot(v[i], u), sz[u] += sz[v[i]], chkmax(maxch[u], sz[v[i]]);
    chkmax(maxch[u], nodesum - sz[u]);
    if (maxch[u] < maxch[rt]) rt = u;
}

void GetDepth(int u, int fa)
{
    lst[++ cnt] = dep[u];
    for (register int i = beg[u]; i; i = nex[i]) if (fa != v[i] && !vis[v[i]])
        dep[v[i]] = dep[u] + w[i], GetDepth(v[i], u);
}

int calc(int u, int init)
{
    dep[u] = init,
    cnt = 0, GetDepth(u, u),
    sort(lst + 1, lst + 1 + cnt);
    register int  l = 1, r = cnt, Sum = 0;
    while (l < r) if (lst[l] + lst[r] <= k) Sum += r - l, ++ l; else -- r;
    return Sum;
}

void Solve(int u)
{
    Ans += calc(u, 0), vis[u] = 1;
    for (register int i = beg[u]; i; i = nex[i]) if (!vis[v[i]]) {
        Ans -= calc(v[i], w[i]);
        nodesum = sz[v[i]], rt = 0, GetRoot(v[i], v[i]);
        Solve(rt);
    }
}

int main()
{
    File();
    register int uu, vv, ww;
    maxch[0] = INF;
    while (scanf("%d%d", &n, &k) != EOF && n) {
        Set(beg, 0), e = 0, Set(vis, 0), Ans = 0;
        For(i, 2, n) uu = read(), vv = read(), ww = read(), add(uu, vv, ww), add(vv, uu, ww);
        nodesum = n, rt = 0, GetRoot(1, 1);
        Solve(rt);
        printf("%d\n", Ans);
    }
    return 0;
}
//竹外桃花三两枝，春江水暖鸭先知。
//    -- 苏轼《惠崇春江晚景》