Tree

题目链接
数据结构/点分治

Description

求树上长度不超过 $k$ 的路径数。

Solution

若节点 $root$ 为根，则对 $root$ 而言，树上路径可分为2类

经过根节点 $root$ 。(分为 $x->root$ , $root->y$ )
包含于 root 的某一棵子树中。(不经过根节点 $root$ )

对于当前节点 $root$ ，dfs 求出子树节点到其的距离 $dis_{node}$ 。

inline void getDis(int cur, int fa) { //求出到根root的距离
    vec.emplace_back(dis[cur]);
    for (auto &e : G[cur]) {
        int to = e.first, w = e.second;
        if (to == fa || vis[to]) continue;
        dis[to] = dis[cur] + w;
        getDis(to, cur);
    }
}

将 $dis[]$ 数组升序排序，使用两个指针 $l, r$ 从前、后开始扫描数组。

容易发现，在 l 从左到右扫描的过程中，恰好使得 $dis[l] + dis[r] \le\ k$ 的 r 是从右向左递减的，于是符合条件的点对有 $r - l$ 对。

inline void calc(int cur, int w, int type) { //cur为有根树的根节点(重心) type为1表示加
    dis[cur] = w, vec.clear(), getDis(cur, 0);
    sort(vec.begin(), vec.end());
    int l = 0, r = (int)vec.size() - 1;
    while (l <= r) {
        if (vec[l] + vec[r] <= k) res += type * (r - l), ++l;
        else --r;
    }
}

不过一些点对在 $root$ 的某一棵子树中，且会使 $root$ 与其儿子间的边被经过两次。可以以 $root$ 的儿子节点为根, $dis[son] = w_{root, son}$ ，再次求 $dis[]$ 并统计, 减去这部分的结果即可 ( $son$ 子树满足条件的在 $root$ 中不满足)。
calc(to, w, -1); //容斥思想减去儿子中的路径数注意dis要加上w

于是可不断选取 $root$ 并统计：
1.若递归深度为 T 层，则总时间复杂度 $O(TNlogN)$ 。
2.若选取的 $root$ 为子树重心，则 $root$ 的子树大小不会超过原子树大小的一半，点分治最多 $logN$ 层。

inline void getRoot(int cur, int fa) { //求树的重心 自上而下
    sz[cur] = 1, dp[cur] = 0;
    for (auto &e : G[cur]) {
        int to = e.first, w = e.second;
        if (to == fa || vis[to]) continue;
        getRoot(to, cur);
        sz[cur] += sz[to];
        dp[cur] = max(dp[cur], sz[to]); //dp[cur]求以cur为根的最大子树大小
    }
    dp[cur] = max(dp[cur], tot - sz[cur]);
    if (dp[cur] < dp[root]) root = cur;
}

时间复杂度 $(NlogNlogN)$ 。

Code

Tree

#include<bits/stdc++.h>

using namespace std;

constexpr int MAXN = 4e4 + 4;

template<typename T> inline void read(T &x) {
    int f = 1; x = 0;
    char ch = getchar();
    while (!isdigit(ch)) f = (ch == '-') ? -1 : 1, ch = getchar();
    while (isdigit(ch)) x = x * 10 + ch - 48, ch = getchar();
    x *= f;
}

vector< pair<int, int> > G[MAXN];
vector<int> vec;

int n, k, tot, root, res;
int sz[MAXN], dp[MAXN], dis[MAXN], vis[MAXN];

inline void getRoot(int cur, int fa) { //求树的重心 自上而下
    sz[cur] = 1, dp[cur] = 0;
    for (auto &e : G[cur]) {
        int to = e.first, w = e.second;
        if (to == fa || vis[to]) continue;
        getRoot(to, cur);
        sz[cur] += sz[to];
        dp[cur] = max(dp[cur], sz[to]); //dp[cur]求以cur为根的最大子树大小
    }
    dp[cur] = max(dp[cur], tot - sz[cur]);
    if (dp[cur] < dp[root]) root = cur;
}

inline void getDis(int cur, int fa) { //求出到根root的距离
    vec.emplace_back(dis[cur]);
    for (auto &e : G[cur]) {
        int to = e.first, w = e.second;
        if (to == fa || vis[to]) continue;
        dis[to] = dis[cur] + w;
        getDis(to, cur);
    }
}

inline void calc(int cur, int w, int type) { //cur为有根树的根节点(重心) type为1表示加
    dis[cur] = w, vec.clear(), getDis(cur, 0);
    sort(vec.begin(), vec.end());
    int l = 0, r = (int)vec.size() - 1;
    while (l <= r) {
        if (vec[l] + vec[r] <= k) res += type * (r - l), ++l;
        else --r;
    }
}

inline void dfs(int cur) {
    vis[cur] = true, calc(root, 0, 1);
    for (auto &e : G[cur]) {
        int to = e.first, w = e.second;
        if (vis[to]) continue;
        calc(to, w, -1); //容斥思想 减去儿子中的路径数 注意要加上w
        root = cur, tot = dp[root] = sz[to]; //准备重新选择根节点
        getRoot(to, root), dfs(root);
    }
}

int main() {
    read(n);
    for (int i = 1, u, v, w; i < n; i++) {
        read(u), read(v), read(w);
        G[u].emplace_back(make_pair(v, w)), G[v].emplace_back(make_pair(u, w));
    }
    read(k);
    root = 0, tot = dp[root] = n;
    getRoot(1, root), dfs(root);
    cout << res << endl;
    return 0;
}