【洛谷P4178】Tree

题面

题解

感觉和\(CDQ\)分治一样套路啊

首先，构建出点分树

对于每一层分治重心，求出它到子树中任意点的距离

然后\(two-pointers\)计算满足小于等于\(K\)的点对数目，加入答案

但是可能会算重，那么就减去子树内两两点之间的贡献即可。

代码

#include<cstdio>
#include<cstring>
#include<algorithm>
#define RG register
#define file(x) freopen(#x".in", "r", stdin);freopen(#x".out", "w", stdout);
#define clear(x, y) memset(x, y, sizeof(x));

namespace IO
{
    const int BUFSIZE = 1 << 20;
    char ibuf[BUFSIZE], *is = ibuf, *it = ibuf;
    inline char getchar() { if (is == it) it = (is = ibuf) + fread(ibuf, 1, BUFSIZE, stdin); return *is++; }
}

inline int read()
{
    int data = 0, w = 1;
    char ch = IO::getchar();
    while(ch != '-' && (ch < '0' || ch > '9')) ch = IO::getchar();
    if(ch == '-') w = -1, ch = IO::getchar();
    while(ch >= '0' && ch <= '9') data = data * 10 + (ch ^ 48), ch = IO::getchar();
    return data*w;
}

const int maxn(40010);
struct edge { int next, to, dis; } e[maxn << 1];
int head[maxn], e_num, n, Size, Max, root, stk[maxn], dep[maxn], top, size[maxn], K, ans, vis[maxn];
inline void add_edge(int from, int to, int dis) { e[++e_num] = (edge) {head[from], to, dis}; head[from] = e_num; }

void GetRoot(int x, int fa)
{
    size[x] = 1; int max = 0;
    for(RG int i = head[x]; i; i = e[i].next)
    {
        int to = e[i].to; if(to == fa || vis[to]) continue;
        GetRoot(to, x); size[x] += size[to]; max = std::max(max, size[to]);
    }
    max = std::max(max, Size - size[x]);
    if(max < Max) Max = max, root = x;
}

void GetDep(int x, int fa)
{
    stk[++top] = dep[x];
    for(RG int i = head[x]; i; i = e[i].next)
    {
        int to = e[i].to; if(to == fa || vis[to]) continue;
        dep[to] = dep[x] + e[i].dis; GetDep(to, x);
    }
}

int Calc(int x, int pre)
{
    top = 0; dep[x] = pre; GetDep(x, 0); std::sort(stk + 1, stk + top + 1);
    int l = 1, r = top, sum = 0;
    while(l < r) { if(stk[l] + stk[r] <= K) sum += r - l, ++l; else --r; }
    return sum;
}

void Solve(int x)
{
    ans += Calc(x, 0); vis[x] = 1;
    for(RG int i = head[x]; i; i = e[i].next)
    {
        int to = e[i].to; if(vis[to]) continue;
        ans -= Calc(to, e[i].dis); Size = Max = size[to];
        GetRoot(to, x); Solve(root);
    }
}

int main()
{
#ifndef ONLINE_JUDGE
    file(cpp);
#endif

    Max = Size = n = read();
    for(RG int i = 1, a, b, c; i < n; i++)
        a = read(), b = read(), c = read(), add_edge(a, b, c), add_edge(b, a, c);
    K = read(); GetRoot(1, 0); Solve(root); printf("%d\n", ans);
    return 0;
}