洛谷 P3806 [模板] 点分治 1 题解

【模板】点分治 1

题目描述

给定一棵有 $n$ 个点的树，询问树上距离为 $k$ 的点对是否存在。

输入格式

第一行两个数 $n,m$。

第 $2$ 到第 $n$ 行，每行三个整数 $u,v,w$，代表树上存在一条连接 $u$ 和 $v$ 边权为 $w$ 的路径。

接下来 $m$ 行，每行一个整数 $k$，代表一次询问。

输出格式

对于每次询问输出一行一个字符串代表答案，存在输出 AYE，否则输出 NAY。

样例 #1

样例输入 #1

2 1
1 2 2
2

样例输出 #1

AYE

数据规模与约定

• 对于 $30\%$ 的数据，保证 $n\le 100$。
• 对于 $60\%$ 的数据，保证 $n\le 1000$，$m\le 50$ 。
• 对于 $100\%$ 的数据，保证 $1\le n\le 10^4$，$1\le m\le 100$，$1\le k\le 10^7$，$1\le u,v\le n$，$1\le w\le 10^4$。

思路

树上的路径分为两类， 经过根节点root的路径和包含于root的某棵子树里(即不经过root)的路径。对于经过根节点的路径，节点u到节点v的路径长度即为u到根节点的路径长度加上v到根节点的路径长度。而对于不经过根节点的路径，可以找到所在子树的根，然后转化为求经过根节点的路径，也就是分治做法。但是需要注意的是，随便找一个根节点进行点分治的做法最坏时间复杂度会达到O($n^2$)，不能接受。因此可以先找到树的重心，以重心作为根节点进行点分治，这样复杂度可以降到O($nlogn$)。此外，对于m次询问的处理，可以先将所有询问读入数组，然后在进行点分治的过程中处理所有询问，判断是否存在距离为k的点对。

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int maxn = 1e4 + 6;
const int maxm = 1e2 + 6;
const int maxk = 1e7 + 6;

int n, m;

struct edge
{
    int to; // 边的指向
    int w;  // 边权
};

vector<edge> e[maxn];

int _size[maxn];  // 以i为根的子树的结点数
int max_size;     // 求重心时的删除某结点后剩余子树结点个数的最大值的最小值
int sum;          // 当前子树的结点数
int root;         // 当前子树的根
int del[maxn];    // 处理过的子树的根标记为已删除
int dis[maxn];    // 以某个点为根时各个点到根节点的距离
bool judge[maxk]; // 到根节点的距离为i是否存在
int query[maxm];  // 记录询问
bool ans[maxm];   // 第i个询问是否可满足

void add_edge(int u, int v, int w) // 加边函数
{
    e[u].push_back({v, w});
}

void get_root(int u, int vis) // 找到重心，以重心为根
{
    int cur_max_num = 0;
    _size[u] = 1;
    for (int i = 0; i < e[u].size(); i++)
    {
        int v = e[u][i].to;
        if (v == vis || del[v])
            continue;
        get_root(v, u);
        _size[u] += _size[v];
        cur_max_num = max(cur_max_num, _size[v]);
    }
    cur_max_num = max(cur_max_num, sum - _size[u]);
    if (cur_max_num < max_size)
    {
        max_size = cur_max_num;
        root = u;
    }
}

void get_distance(int u, int vis, vector<int> &len) // 求子树中各个点到根节点的距离
{
    len.push_back(dis[u]);
    for (int i = 0; i < e[u].size(); i++)
    {
        int v = e[u][i].to;
        if (v == vis || del[v])
            continue;
        dis[v] = dis[u] + e[u][i].w;
        get_distance(v, u, len);
    }
}

void cal_path_length(int u) // 计算路径长度
{
    del[u] = 1;   // 记录以u为根的树已经计算过路径
    judge[0] = 1; // 预处理距离为0存在

    set<int> need_to_clear; // 记录需要清空的judge下标

    for (int i = 0; i < e[u].size(); i++)
    {
        int v = e[u][i].to;
        if (del[v])
            continue;
        // 离散化存储该子树中所有距离根结点的路径长度
        vector<int> len;
        // 求出子树v的各个结点到根u的距离
        dis[v] = e[u][i].w;
        get_distance(v, u, len);
        // 枚举距离和询问，将各个子树到根u的距离进行组合求和，判断是否存在该询问对应的距离之和
        for (int j = 0; j < len.size(); j++)
        {
            for (int k = 1; k <= m; k++)
            {
                if (query[k] >= len[j])
                {
                    if (judge[query[k] - len[j]] == 1)
                        ans[k] = 1;
                }
            }
        }
        // 记录先前子树到根的距离，用于后续子树与先前子树的组合求和
        for (int j = 0; j < len.size(); j++)
        {
            if (len[j] <= 1e7)
            {
                judge[len[j]] = 1;
                need_to_clear.insert(len[j]);
            }
        }
    }
    // 清空judge
    for (set<int>::iterator it = need_to_clear.begin(); it != need_to_clear.end(); it++)
    {
        judge[*it] = 0;
    }
}

void divide(int u) // 分治
{
    cal_path_length(u); // 计算经过根u的所有路径长度
    // 对u的子树分治
    for (int i = 0; i < e[u].size(); i++)
    {
        int v = e[u][i].to;
        if (del[v])
            continue;
        // 对子树进行分治，要更新sum和max_size为子树的大小
        sum = _size[v];
        max_size = _size[v];
        get_root(v, 0);
        divide(root);
    }
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    cin >> n >> m;
    int a, b, c; // 连接a和b的边权为c的路径
    for (int i = 1; i <= n - 1; i++)
    {
        cin >> a >> b >> c;
        add_edge(a, b, c);
        add_edge(b, a, c);
    }

    int k;
    for (int i = 1; i <= m; i++)
    {
        cin >> k;
        query[i] = k;
    }

    sum = n;
    max_size = n;
    get_root(1, 0);    // 寻找重心
    get_root(root, 0); // 通过该重心，更新所有子树的大小
    divide(root);

    for (int i = 1; i <= m; i++)
    {
        if (ans[i])
            cout << "AYE" << '\n';
        else
            cout << "NAY" << '\n';
    }

    return 0;
}