AcWing252. 树（点分治+容斥/树状数组）

AcWing252. 树

题意

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

（本题与洛谷 P4178 Tree 非常相似，但是这题 $k$ 比较大，且存在零边，当然也可以和洛谷那题一样用树状数组+特判零边来完成，这里主要讲一下容斥的方法。）

解法

• 考虑以 $u$ 作为路径的lca， $u$ 连接了若干子树，遍历每一棵子树，计算出深度。将所有的深度排序（包括根），然后用双指针来计算有多少对距离之和小于等于 $k$ 。这样计算的包含一些不合法的路径，可以发现当两个点来自不同的子树的时候一定是合法的，不合法来自于两个点在用一个子树上，这样lca就不对了。所以再减去同一棵子树的贡献即可。
• 每次枚举的lca 都应该是树的重心，否则可以在链上会退化成 $O(n^2)$ 。

代码

容斥：

#pragma region
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <vector>
using namespace std;
typedef long long ll;
#define tr t[root]
#define lson t[root << 1]
#define rson t[root << 1 | 1]
#define rep(i, a, n) for (int i = a; i <= n; ++i)
#define per(i, a, n) for (int i = n; i >= a; --i)
#pragma endregion
const int maxn = 1e4 + 5;
int n, k;
vector<pair<int, int>> g[maxn];
int rt, d[maxn], sz[maxn];
bool vis[maxn];
void dfs_rt(int u, int f, int tot) {
    sz[u] = 1;
    int maxx = 0;
    for (auto e : g[u]) {
        int v = e.first;
        if (vis[v] || v == f) continue;
        dfs_rt(v, u, tot);
        sz[u] += sz[v];
        maxx = max(maxx, sz[v]);
    }
    maxx = max(maxx, tot - sz[u]);
    if (2 * maxx <= tot) rt = u;
}
int cnt, tmp[maxn];
void dfs_d(int u, int f) {
    tmp[++cnt] = d[u];
    for (auto e : g[u]) {
        int v = e.first, w = e.second;
        if (vis[v] || v == f) continue;
        d[v] = d[u] + w;
        dfs_d(v, u);
    }
}
int cul(int u, int f) {
    cnt = 0;
    dfs_d(u, f);
    sort(tmp + 1, tmp + 1 + cnt);
    int l = 1, r = cnt, ans = 0;
    while (l < r) {
        while (l < r - 1 && tmp[l] + tmp[r] > k) --r;
        if (tmp[l] + tmp[r] > k) break;
        ans += r - l;
        ++l;
    }
    return ans;
}
int work(int u, int f, int tot) {
    dfs_rt(u, f, tot);
    u = rt, vis[u] = 1, d[u] = 0;
    int ans = cul(u, f);
    for (auto e : g[u]) {
        int v = e.first;
        if (vis[v]) continue;
        ans -= cul(v, u);
        ans += work(v, u, cnt);
    }
    return ans;
}
int main() {
    while (scanf("%d%d", &n, &k) && n && k) {
        rep(i, 1, n) g[i].clear(), vis[i] = 0;
        rep(i, 1, n - 1) {
            int u, v, w;
            scanf("%d%d%d", &u, &v, &w);
            ++u, ++v;
            g[u].push_back({v, w});
            g[v].push_back({u, w});
        }
        printf("%d\n", work(1, 0, n));
    }
}

树状数组：

#pragma region
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <vector>
using namespace std;
typedef long long ll;
#define tr t[root]
#define lson t[root << 1]
#define rson t[root << 1 | 1]
#define rep(i, a, n) for (int i = a; i <= n; ++i)
#define per(i, a, n) for (int i = n; i >= a; --i)
#pragma endregion
const int maxn = 1e4 + 5;
int n, k;
vector<pair<int, int>> g[maxn];
int sz[maxn], rt, d[maxn];
bool vis[maxn];
int c[int(5e6 + 5)];
void add(int pos, int x) {
    while (pos <= k) c[pos] += x, pos += pos & -pos;
}
int getsum(int pos) {
    int ans = 0;
    while (pos) ans += c[pos], pos -= pos & -pos;
    return ans + c[0];
}
void init() {
    rep(i, 1, n) g[i].clear(), vis[i] = 0;
}
void dfs_rt(int u, int f, int tot) {
    sz[u] = 1;
    int maxx = 0;
    for (auto e : g[u]) {
        int v = e.first;
        if (vis[v] || v == f) continue;
        dfs_rt(v, u, tot);
        sz[u] += sz[v];
        maxx = max(maxx, sz[v]);
    }
    maxx = max(maxx, tot - sz[u]);
    if (maxx * 2 <= tot) rt = u;
}
void dfs_ans(int u, int f, int &ans, int &cnt) {
    ++cnt;
    if (d[u] <= k) ans += getsum(k - d[u]) + 1;
    for (auto e : g[u]) {
        int v = e.first, w = e.second;
        if (vis[v] || v == f) continue;
        d[v] = d[u] + w;
        dfs_ans(v, u, ans, cnt);
    }
}
void dfs_update(int u, int f, int val) {
    if (d[u] <= k) {
        if (!d[u])
            c[d[u]] += val;
        else
            add(d[u], val);
    }
    for (auto e : g[u]) {
        int v = e.first;
        if (vis[v] || v == f) continue;
        dfs_update(v, u, val);
    }
}
int work(int u, int f, int tot) {
    dfs_rt(u, f, tot);
    u = rt, vis[u] = 1, d[u] = 0;
    int ans = 0;
    for (auto e : g[u]) {
        int v = e.first, w = e.second, cnt;
        if (vis[v]) continue;
        d[v] = w, cnt = 0;
        dfs_ans(v, u, ans, cnt);
        sz[v] = cnt;
        dfs_update(v, u, 1);
    }
    for (auto e : g[u]) {
        int v = e.first;
        if (vis[v]) continue;
        dfs_update(v, u, -1);
    }
    for (auto e : g[u]) {
        int v = e.first;
        if (vis[v]) continue;
        ans += work(v, u, sz[v]);
    }
    return ans;
}
int main() {
    while (scanf("%d%d", &n, &k) && n && k) {
        init();
        rep(i, 1, n - 1) {
            int u, v, w;
            scanf("%d%d%d", &u, &v, &w), ++u, ++v;
            g[u].push_back({v, w});
            g[v].push_back({u, w});
        }
        printf("%d\n", work(1, 0, n));
    }
}