P4178 Tree

给定一棵 $n$ 个节点的树，每条边有边权，求出树上两点距离小于等于 $k$ 的点对数量。

考虑答案可能从哪些状态中得到：

两个端点在同一子树内；
两个端点在不同的子树内；
其中一个点是重心。

对于点分治当前的点，维护一个集合 $p$ ，为以该点为根的子树中所有的距离；对于该点的子节点，维护一个集合 $q$ （都为可重集），为以该子节点为根的子树中所有的距离。显然集合 $p$ 是集合是所有子节点集合 $q$ 的交集。这样前两种答案可以表示为 $f(p) - \displaystyle \sum _ {(p, q) \in E} f(q)$ 。

对于一个集合，如何计算答案？可以排序后使用双指针算法。

查看代码

#include <cstdio>
#include <algorithm>
#include <vector>
using namespace std;
const int N = 4e4 + 10, M = 8e4 + 10;
bool vis[N];
int n, m;
int idx, hd[N], nxt[M], edg[M], wt[M];
int Size(int x, int fa)
{
    if (vis[x])
        return 0;
    int res = 1;
    for (int i = hd[x]; ~i; i = nxt[i])
        if (edg[i] != fa)
            res += Size(edg[i], x);
    return res;
}
int WeightCentre(int x, int fa, int tot, int &wc)
{
    if (vis[x])
        return 0;
    int sum = 1, mx = 0;
    for (int i = hd[x]; ~i; i = nxt[i])
        if (edg[i] != fa)
        {
            int t = WeightCentre(edg[i], x, tot, wc);
            mx = max(mx, t);
            sum += t;
        }
    mx = max(mx, tot - sum);
    if (mx <= tot / 2)
        wc = x;
    return sum;
}
void Dist(int x, int fa, int d, vector<int> &q)
{
    if (vis[x])
        return;
    q.push_back(d);
    for (int i = hd[x]; ~i; i = nxt[i])
        if (edg[i] != fa)
            Dist(edg[i], x, d + wt[i], q);
}
int get(vector<int> k)
{
    sort(k.begin(), k.end());
    int res = 0;
    for (int i = k.size() - 1, j = -1; i >= 0; i--)
    {
        while (j + 1 < i && k[j + 1] + k[i] <= m)
            j++;
        j = min(j, i - 1);
        res += j + 1;
    }
    return res;
}
int cal(int x)
{
    if (vis[x])
        return 0;
    WeightCentre(x, -1, Size(x, -1), x);
    vis[x] = true;
    int res = 0;
    vector<int> p;
    for (int i = hd[x]; ~i; i = nxt[i])
    {
        vector<int> q;
        Dist(edg[i], -1, wt[i], q);
        res -= get(q);
        for (int k : q)
        {
            if (k <= m)
                res++;
            p.push_back(k);
        }
    }
    res += get(p);
    for (int i = hd[x]; ~i; i = nxt[i])
        res += cal(edg[i]);
    return res;
}
void add(int a, int b, int c)
{
    nxt[++idx] = hd[a];
    hd[a] = idx;
    edg[idx] = b;
    wt[idx] = c;
}
int main()
{
    scanf("%d", &n);
    for (int i = 1; i <= n; i++)
        hd[i] = -1;
    for (int i = 1, a, b, c; i < n; i++)
    {
        scanf("%d%d%d", &a, &b, &c);
        add(a, b, c);
        add(b, a, c);
    }
    scanf("%d", &m);
    printf("%d", cal(1));
    return 0;
}