CodeForces 161D Distance in Tree 点分治-蒲公英云

CodeForces 161D Distance in Tree 点分治

忘是亡心i 2023-05-21 07:23 120阅读 0赞

题目链接: http://codeforces.com/problemset/problem/161/D
题意：给你一棵树，让你求有多少对点的距离等于k 点分治裸题，每次分治子树记录经过当前根节点的距离

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
#define fi first
#define se second
#define ls rt << 1
#define rs rt << 1|1
#define po pop_back
#define pb push_back
#define mk make_pair
#define lson l, mid, ls
#define rson mid + 1, r, rs
#define pll pair<ll, ll>
#define pii pair<int, int>
#define ull unsigned long long
#define pdd pair<double, double>
const int mod = 1e9 + 7;
const int maxn = 5e4 + 10;
const int inf = 0x3f3f3f3f;
const ll linf = 0x3f3f3f3f3f3f3f3f;
int head[maxn * 2], nxt[maxn * 2], to[maxn * 2], ver[maxn * 2], tot1 = -1;
void add(int u, int v, int w)
{
    nxt[++tot1] = head[u]; //当前边的后继
    head[u] = tot1; //起点u的第一条边
    to[tot1] = v; //当前边的终点
    ver[tot1] = w; //当前边的权值
}
vector<int> vc; //离线询问
int sz[maxn], Ma[maxn], vis[maxn]; //sz子树大小 Ma[u]删除u的子树中最大的树大小
int dis[maxn], q[maxn], id[maxn]; //q记录处理子树时出现过的距离值 id记录q方便清空res
int rt, tot, R; //分治根 子树总结点数 记录长度
ll res[1000010], ans[1000010]; //从分治点出现的路径长度  所有出现过的路径长度
void getrt(int u, int fa) //得到子树的非带权重心
{
    sz[u] = 1, Ma[u] = 0; //初始化
    for(int i = head[u]; ~i; i = nxt[i])
    {
        int v = to[i];
        if(vis[v] || v == fa) continue;
        getrt(v, u); //递归得到子树大小
        sz[u] += sz[v];
        Ma[u] = max(Ma[u], sz[v]); //更新u结点的Ma
    }
    Ma[u] = max(Ma[u], tot - sz[u]);
    if(Ma[u] < Ma[rt]) rt = u; //更新当前子树的重心
}
void dfs(int u, int fa)
{
    q[++R] = dis[u]; //记录所有出现的距离值
    for(int i = head[u]; ~i; i = nxt[i])
    {
        int v = to[i], w = ver[i];
        if(vis[v] || v == fa) continue;
        dis[v] = dis[u] + w;
        //if(dis[v] > 10000000) continue;
        dfs(v, u); //可以进行剪枝例如大于1e7的剪掉
    }
}
void calc(int u)
{
    int pos = 0;
    for(int i = head[u]; ~i; i = nxt[i])
    {
        int v = to[i], w = ver[i];
        if(vis[v]) continue;
        R = 0; //初始化
        dis[v] = w; //第一个点的距离
        dfs(v, u); //处理每个子树的距离
        for(int k = R; k; --k) //遍历当前子树的距离
        {
            for(int j = 0; j < vc.size(); ++j)
                if(vc[j] >= q[k]) //对于每个询问如果路径存在就标记询问
                    ans[vc[j]] += res[vc[j] - q[k]];
        }
        for(int k = R; k; --k) //标记出现过的距离
            res[q[k]]++, id[++pos] = q[k];
    }
    for(int i = 1; i <= pos; ++i) //每处理完一个子树就清空
        res[id[i]] = 0;
}
void slove(int u)
{
    vis[u] = 1, res[0] = 1; //res[i]表示到根距离为i的路径是否存在
    calc(u); //处理以u为根的子树
    for(int i = head[u]; ~i; i = nxt[i]) //对每个子树进行分治
    {
        int v = to[i];
        if(vis[v]) continue;
        tot = sz[v], Ma[rt = 0] = inf; //设置子树大小
        getrt(v, 0); //找到新的rt
        slove(rt);
    }
}
int main()
{
    memset(head, -1, sizeof(head));
    int n, k;
    scanf("%d%d", &n, &k);
    for(int i = 1; i < n; ++i)
    {
        int u, v, w = 1;
        scanf("%d%d", &u, &v);
        add(u, v, w), add(v, u, w);
    }
    vc.pb(k);
    tot = Ma[rt] = n; //首次先找整棵树的重心
    getrt(1, 0);
    slove(rt);
    printf("%lld\n", ans[k]);
    return 0;
}