洛谷P4178 Tree 題解 樹上點分治

題目鏈接：https://www.luogu.com.cn/problem/P4178

解題思路：

點分治模板題。

設當前重心爲 \(u\)，一共有三種不同類型的路徑：

1. 路徑的一個端點恰好是重心 \(u\)；
2. 路徑的兩個端點在 \(u\) 的不同子樹中；
3. 路徑的兩個端點在 \(u\) 的同一個子樹中。

找到重心 \(u\) 之後，前兩類路徑分開求。第三類路徑分治求。

示例程序：

#include <bits/stdc++.h>
using namespace std;
const int maxn = 4e4 + 5;

struct Edge {
    int v, w;
};
vector<Edge> g[maxn];
int n, K, ans;
bool vis[maxn];

int get_size(int u, int p) {
    if (vis[u]) return 0;
    int sum = 1;
    for (auto e : g[u])
        if (e.v != p)
            sum += get_size(e.v, u);
    return sum;
}

int get_wc(int u, int p, int tot, int &wc) {
    if (vis[u]) return 0;
    int sum = 1, ms = 0;
    for (auto e : g[u]) {
        int v = e.v;
        if (v == p) continue;
        int tmp = get_wc(v, u, tot, wc);
        ms = max(ms, tmp);
        sum += tmp;
    }
    ms = max(ms, tot - sum);
    if (ms <= tot / 2) wc = u;
    return sum;
}

void get_dfs(int u, int p, int dis, vector<int> &v) {
    if (vis[u] || dis > K) return;
    v.push_back(dis);
    for (auto e : g[u])
        if (e.v != p)
            get_dfs(e.v, u, dis + e.w, v);
}

void dfs(int u) {
    if (vis[u]) return;
    get_wc(u, -1, get_size(u, -1), u);
    vis[u] = true;
    vector<int> v1, v2;
    for (auto e : g[u])
        get_dfs(e.v, u, e.w, v1);
    sort(v1.begin(), v1.end());
    int tmp = 0;
    for (auto e : g[u]) {
        v2.clear();
        get_dfs(e.v, u, e.w, v2);
        sort(v2.begin(), v2.end());
        for (auto d : v2) {
            int num1 = upper_bound(v1.begin(), v1.end(), K - d) - v1.begin();
            int num2 = upper_bound(v2.begin(), v2.end(), K - d) - v2.begin();
            tmp += num1 - num2;
        }
    }
    ans += v1.size();
    ans += tmp / 2;

    for (auto e : g[u])
        dfs(e.v);
}

int main() {
    scanf("%d", &n);
    for (int i = 1; i < n; i++) {
        int u, v, w;
        scanf("%d%d%d", &u, &v, &w);
        g[u].push_back({v, w});
        g[v].push_back({u, w});
    }
    scanf("%d", &K);
    dfs(1);
    printf("%d\n", ans);
    return 0;
}