AcWing - 252 - 树 = 点分治

https://www.acwing.com/problem/content/description/254/

这个才是真的点分治，上次那个直接树dp就结束了，因为和子树的size没有什么关系，甚至那个重心是自己骗自己的。

这个点分就厉害了，每次都要找新的重心。

抄别人代码：

#include <map>
#include <set>
#include <list>
#include <cmath>
#include <deque>
#include <queue>
#include <stack>
#include <cctype>
#include <string>
#include <cstdio>
#include <vector>
#include <complex>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <algorithm>
#include <functional>
using namespace std;
typedef long long LL;
const int maxn = 40000 + 100;

template <class T>
inline void read(T &s) {
    s = 0;
    T w = 1, ch = getchar();
    while (!isdigit(ch)) { if (ch == '-') w = -1; ch = getchar(); }
    while (isdigit(ch)) { s = (s << 1) + (s << 3) + (ch ^ 48); ch = getchar(); }
    s *= w;
}

int n, m, k, tot, all_node, ans, root;
// root 为重心
// ans 为答案
// all_node 表示所有点的数目
int lin[maxn], max_part[maxn], size[maxn], len[maxn], d[maxn];
// d[i] 表示i到根节点的距离
// len[i] 表示路径长度
// len[0] 中记录节点个数
// size[i] 表示以i为根的子树大小
// max_part[i] 表示i的最大的子树大小
bool vis[maxn];
// vis用于求重心时避免重复访问
struct node {
    int next, to, dis;
} edge[maxn << 1];
// edge 用于存边

inline void add(int from, int to, int dis) {
    edge[++tot].to = to;
    edge[tot].dis = dis;
    edge[tot].next = lin[from];
    lin[from] = tot;
}

void get_root(int u, int fa) { // 找到树的重心
    max_part[u] = 0, size[u] = 1;
    for (int i = lin[u]; i; i = edge[i].next) {
        int v = edge[i].to;
        if (v == fa || vis[v]) continue;
        get_root(v, u);
        size[u] += size[v];
        max_part[u] = max(max_part[u], size[v]);
    }
    max_part[u] = max(max_part[u], all_node - size[u]);
    if (max_part[u] < max_part[root]) root = u;
}

void get_dis(int u, int fa) { // 求出每个点到根节点的距离
    len[++len[0]] = d[u];
    for (int i = lin[u]; i; i = edge[i].next) {
        int v = edge[i].to;
        if (v == fa || vis[v]) continue;
        d[v] = d[u] + edge[i].dis;
        get_dis(v, u);
    }
}

int cal(int u, int now) { // 计算以u为根的所有情况的答案
    d[u] = now, len[0] = 0;
    get_dis(u, 0);
    sort(len + 1, len + len[0] + 1);
    int all = 0;
    for (int l = 1, r = len[0]; l < r; ) {
        if (len[l] + len[r] <= k) {
            all += r - l;
            ++l;
        }
        else r--;
    }
    return all;
}

void solve(int u) { // 求解以u为重心的情况
    vis[u] = true;
    ans += cal(u, 0);
    for (int i = lin[u]; i; i = edge[i].next) {
        int v = edge[i].to;
        if (vis[v]) continue;
        ans -= cal(v, edge[i].dis); // 减去不合法的
        all_node = size[v];
        root = 0;
        get_root(v, u);
        solve(root);
    }
}

int main() {
    while (1) {
        read(n), read(k);
        if (!n && !k) break;

        tot = 0; ans = 0;
        memset(lin, 0, sizeof(lin));
        memset(max_part, 0, sizeof(max_part));
        memset(size, 0, sizeof(size));
        memset(vis, false, sizeof(vis));

        for (int i = 1; i < n; ++i) {
            int u, v, w;
            read(u), read(v), read(w);
            add(u+1, v+1, w);
            add(v+1, u+1, w);
        }
        // read(k);
        all_node = n, max_part[0] = n, root = 0;
        get_root(1, 0);
        solve(root);
        printf("%d\n", ans);
    }
    return 0;
} 

作者：MILLOPE
链接：https://www.acwing.com/solution/AcWing/content/2826/
来源：AcWing
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。