Tree

题目描述

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

输入格式

第一行输入一个整数 $n$，表示节点个数。

第二行到第 $n$ 行每行输入三个整数 $u,v,w$ ，表示 $u$ 与 $v$ 有一条边，边权是 $w$。

第 $n+1$ 行一个整数 $k$ 。

输出格式

一行一个整数，表示答案。

样例 #1

样例输入 #1

7
1 6 13 
6 3 9 
3 5 7 
4 1 3 
2 4 20 
4 7 2 
10

样例输出 #1

5

数据规模与约定

对于全部的测试点，保证：

• $1\le n\le 4\times 10^4$。
• $1\le u,v\le n$。
• $0\le w\le 10^3$。
• $0\le k\le 2\times 10^4$。

原题

洛谷P4178——传送门

思路

相较于点分治模板题洛谷P3806，本题要求求出距离小于等于 $k$ 的点对数量。我的做法是点分治+容斥：对于以某一个重心为根的树中点对的统计，采用将所有点到根节点的距离进行排序，然后用双指针求出距离相加小于等于 $k$ 的点对数量，并减去其中两点在同一个子树中的情况(即减去以子节点为根的子树中的点对数量)。

代码

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

const int maxn = 4e4 + 6;
const int maxk = 2e4 + 6;

int n;

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];   // 以某个点为根时各个点到根节点的距离
int k;
int ans = 0;

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);
    }
}

int calc_pair_num(vector<int> &a) // 计算点的对数
{
    sort(a.begin(), a.end());
    int res = 0;
    for (int i = 0, j = a.size() - 1; i < a.size(); i++)
    {
        while (j >= 0 && a[i] + a[j] > k)
        {
            j--;
        }
        if (j >= i)
            res += j;
        else
            res += j + 1;
    }
    return res / 2; // 由于双指针统计的时候重复计算点对(a,b)和(b,a)，所以需要除以2
}

void count_path(int u) // 统计点对数
{
    del[u] = 1; // 记录以u为根的树已经计算过路径

    vector<int> num;
    num.push_back(0);

    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);

        ans -= calc_pair_num(len); // 容斥，减去各个子树中的点对个数
        num.insert(num.begin(), len.begin(), len.end());
    }

    ans += calc_pair_num(num); // 加上该树中的点对个数
}

void divide(int u) // 分治
{
    count_path(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;
    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);
    }

    cin >> k;

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

    cout << ans << '\n';

    return 0;
}