CodeForces 161D Distance in Tree 点分治

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

题目链接: [http://codeforces.com/problemset/problem/161/D][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;
    }

[http_codeforces.com_problemset_problem_161_D]: http://codeforces.com/problemset/problem/161/D