点分治 笔记

本文原在 2024-07-26 09:12 发布于本人洛谷博客。

一、介绍

点分治常用于解决树上路径数目，符合条件的点对数目的问题。

二、实现

先看例题：

P4178 Tree

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

0. 大体思路

找到一棵树的重心，然后统计经过重心的路径数目，然后再分治所有的子树。

由于找的都是重心，所以时间复杂度 \(O(n\log n)\)。

1. 找根

由于处理重心，所以要先写一个函数找重心。

我们知道，以重心为根，重心最大的子树的大小，是绝对小于以非重心为根，那个点的最大子树大小的，所以我们只需要处理每个点的子树大小 \(sz\)，和该点最大的子树的大小 \(mx\) 即可。

需要注意的是，节点 \(u\) 的上方也是一棵子树，而该子树的大小为 \(nodeCnt-sz_u\)，其中 \(nodeCnt\) 表示整棵树的大小。

Code：

void getrt(int u, int fa, int nodecnt) {
    sz[u] = 1;
    mx[u] = 0;
    for (int i = head[u]; i; i = edge[i].next) {
        int v = edge[i].v;
        if (v == fa or vis[v]) continue;
        getrt(v, u, nodecnt);
        sz[u] += sz[v];
        mx[u] = max(mx[u], sz[v]);
    }
    mx[u] = max(mx[u], nodecnt - sz[u]);
    if (mx[u] < mx[rt]) rt = u;
}

2. 处理每个点到重心的距离

既然要处理经过重心，且距离小于等于 \(k\) 的路径，则必须要先求出每个点到根节点的路径，深搜遍历一次即可。

Code：

void getdis(int u, int fa, int dis) {
    q[++p] = dis;
    for (int i = head[u]; i; i = edge[i].next) {
        int v = edge[i].v;
        if (v == fa or vis[v]) continue;
        getdis(v, u, dis + edge[i].w);
    }
}

3. 计算+递归

在上一步中，我们把每个点到重心的距离存入了一个 \(q\) 数组，如果某个点到重心的距离为 \(q_i\)，则答案明显为 \(q\) 数组中，小于等于 \(k-q_i\) 的数的数量。

但是，我们发现这样会算重：

如上图所示，如果 \(k=6\)，则 \(3-2-1-2-4\) 会被计算一次，而 \(3-2-4\) 时又会计算一次。

我们发现，这些重复的路径都满足经过他的子节点 \(2\)，所以减去经过他的子节点的即可。

Code：

int calc() {
    int ret = 0;
    sort(q + 1, q + p + 1);
    for (int i = 1; i < p; i++)
        ret += upper_bound(q + i + 1, q + p + 1, k - q[i]) - (q + i + 1);
    return ret;
}
void solve(int u, int nodecnt) {
    mx[rt = 0] = oo;
    getrt(u, 0, nodecnt);
    getrt(rt, 0, nodecnt);
    p = 0;
    getdis(rt, 0, 0);
    ans += calc();
    vis[rt] = 1;
    for (int i = head[rt]; i; i = edge[i].next) {
        int v = edge[i].v;
        if (vis[v]) continue;
        p = 0;
        getdis(v, rt, edge[i].w);
        ans -= calc();
        solve(v, sz[v]);
    }
}