P4178 Tree （点分治）

题目链接

一：我们考虑树上两点之间的路径有什么情况
1：经过根节点（即在根节点的两端）
2：不经过根节点（完全在一颗子树的一侧）

二：我们考虑这两种路径是否可以归为一类
1：对于第一种的情况两点间路径长度dis( u , v ) 可以看为dis ( u , root ) + dis ( root , v ) 
2：而对于第二种情况我们可以找到u , v 路径所在的子树，将根节点变为子树的根节点，然后就变为了第一种情况
3：总之，所以我们可以不断的递归根节点将所有情况转化为第一种情况
设b[ u ] 表示u属于哪一棵子树(b数组在算法进阶指南里说了，但是代码好像没有体现，而是用了容斥原理）
那么b[ u ] = b[ v ] ，d[ u ] + d[ v ] ≤ k满足条件的第一类路径条数即为满足条件的点对数

三：下面我们来讨论如何计算满足条件的点对数：
1：将目前子树中的节点到根节点的权值放入数组len中，并排序
用两个指针 L,R 扫描数组，每找到一个合法的答案就加上R −L ，就让L++，否则让 R --
2：发现问题

在指针扫描的过程中我们会出现不合的情况

根据len的定义我们存放的一颗树里面所有的边，这是我们计算calc计算时，可能会加上两条边在子树的同一侧的情况，我们可以发现不合法的路径在当前根的一颗子树上，即没有跨越两棵子树（如果跨越了它就合法了），所以们可以在遍历的时候减去不合法的情况（容斥）

3：具体的方法
ans+ = calc( u , 0 ) 
ans − =calc ( v , dis ( u , v ) ) 

这样即可做到不重不漏

四：总结

点分治步骤

• 找到树的重心
• 将重心视为根节点，那么树上任意两点有两种情况
• 路径经过根节点
• 路径不经过根节点
• 通过 calc  函数计算出第一种情况下的答案，把根节点从树中删去
• 对每棵子树递归执行上面的操作

calc函数的计算方法

• 计算出每个结点到根节点的距离 d [ i ]
• 将树上的结点按照 d [ i ] 递增排序
• 指针 l  指向 d [ 1 ] ，指针 r 指向 d [ n ]。
• 若 l 与 r 指向结点的距离小于 k  ,则 ans + = r − l + 1 , l + + 。
• 否则 r − −。
• 当 l > = r  的时候退出循环。
• 按照上面的方法，会把不经过根节点的路径也算入进去。利用容斥原理修正答案：

五：更具体的做法，见注释

#include <bits/stdc++.h>
using namespace std;
#define pi acos(-1)
#define xx first
#define yy second
#define endl "\n"
#define lowbit(x) x & (-x)
#define int long long
#define ull unsigned long long
#define pb push_back
typedef pair<int, int> PII;
typedef pair<double, double> PDD;
#define max(a, b) (((a) > (b)) ? (a) : (b))
#define min(a, b) (((a) < (b)) ? (a) : (b))
#define Ysanqian ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
const int N = 1e6 + 10, M = 1010, inf = 0x3f3f3f3f, mod = 18446, P = 13331;
const double eps = 1e-8;
int n, k, root, sum, ans;
vector<PII> g[N];
int siz[N], max_part[N], d[N], b[N], len[N];
bool st[N];
void get_root(int u, int fa) // 找出重心
{
    siz[u] = 1, max_part[u] = 0; // siz[i]为其子树的大小，max_part[i]为去掉i后的最大连通块大小
    for (auto ed : g[u])
    {
        int v = ed.xx;
        if (v == fa || st[v])
            continue;
        get_root(v, u);
        siz[u] += siz[v];
        max_part[u] = max(max_part[u], siz[v]);
    }
    max_part[u] = max(max_part[u], sum - siz[u]);
    if (max_part[u] < max_part[root]) // 一开始root为0，且将其赋值为 n即为max
        root = u;
}
void get_dis(int u, int fa) // 找出这颗子树的距离。
{
    len[++len[0]] = d[u]; // len[0]存放的是子树节点个数，len[i] 表示子树中的路径长度
    for (auto ed : g[u])
    {
        int v = ed.xx, w = ed.yy;
        if (v == fa || st[v])
            continue;
        d[v] = d[u] + w; // d[i] 表示i到根节点的距离，由上而下的去跟新d数组，因为子节点距离有父节点跟新而来，故在计算距离时只更新父节点即可
        get_dis(v, u);
    }
}
int calc(int u, int w) // 计算以u为根的第一种情况的答案，也就是在树的两边的，但是对于只在一边得还未计算
{
    d[u] = w, len[0] = 0;            // 先给根节点一个值，如果为0说明在递归子树，不是0那就是再利用容斥出去冗余，将该子树的边数先归零
    get_dis(u, 0);                   // 计算改子树的距离根节点的距离
    sort(len + 1, len + len[0] + 1); // 将距离从小到大排序
    int res = 0;
    for (int l = 1, r = len[0]; l < r;) // 直到不论左移还是右移都大于k,if和else if 都进不去
    {
        if (len[l] + len[r] <= k) // l与最大的相加都小于k，其余也小于k
        {
            res += r - l; // 故加上r-l
            ++l;          // 右移左端点
        }
        else
            r--; // 否则左移右端点
    }
    return res; // 返回以u为根的所有情况的答案
}
void solve(int u)
{
    st[u] = 1;           // 标记删除的重心，防止重复计算
    ans += calc(u, 0);   // 加上以我为跟答案
    for (auto ed : g[u]) // 遍历子树
    {
        int v = ed.xx, w = ed.yy;
        if (st[v])
            continue;
        ans -= calc(v, w); // 减去不合法得情况，不合法的情况全都发生在一课子树上，这里为什么传w，因为我们要加上u与v的边权
        sum = siz[v];      // 不要忘记在以v为根节点时，要修改all_node 的大小
        root = 0;
        get_root(v, u); // 递归根节点，但是肯定会有重复
        solve(root);    // 递归解决
    }
}
signed main()
{
    Ysanqian;
    cin >> n;
    sum = n;
    max_part[0] = n; // 让max_part[0]为n当一个哨兵，以便跟新max_part[i]
    for (int i = 1; i < n; i++)
    {
        int a, b, c;
        cin >> a >> b >> c;
        g[a].pb({b, c}); // 建边
        g[b].pb({a, c});
    }
    cin >> k;
    get_root(1, -1); // 先找重心
    solve(root);     //
    cout << ans << endl;
    return 0;
}