洛谷P2634 [国家集训队]聪聪可可点分治+容斥

电玩女神 2023-05-21 07:23 16阅读 0赞

题目链接：[https://www.luogu.com.cn/problem/P2634][https_www.luogu.com.cn_problem_P2634]  
题意：问你有多少点对之间的距离是3的倍数  
思路：点分治，每次分治根u时统计经过u的距离%3，那么当前计算的答案就是res\[0\] \* res\[0\] + res\[1\] \* res\[2\] \* 2,   
可以找样例算一下，然后算根的时候要容斥一下，其实也就是减去多算的儿子答案

#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 = 3e4 + 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[5], ans = 0; //从分治点出现的路径长度  所有出现过的路径长度
    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)
    {
        res[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) % 3;
            //if(dis[v] > 10000000) continue;
            dfs(v, u); //可以进行剪枝例如大于1e7的剪掉
        }
    }
    int calc(int u, int val)
    {
        dis[u] = val % 3;
        for(int i = 0; i <= 2; ++i)
            res[i] = 0;
        dfs(u, 0);
        return res[1] * res[2] * 2 + res[0] * res[0];
    }
    void slove(int u)
    {
        vis[u] = 1, res[0] = 1; //res[i]表示到根距离为i的路径是否存在
        ans += calc(u, 0); //处理以u为根的子树
        for(int i = head[u]; ~i; i = nxt[i]) //对每个子树进行分治
        {
            int v = to[i];
            if(vis[v]) continue;
            ans -= calc(v, ver[i]); //容斥
            tot = sz[v], Ma[rt = 0] = inf; //设置子树大小
            getrt(v, 0); //找到新的rt
            slove(rt);
        }
    }
    
    int main()
    {
        memset(head, -1, sizeof(head));
        int n;
        scanf("%d", &n);
        for(int i = 1; i < n; ++i)
        {
            int u, v, w;
            scanf("%d%d%d", &u, &v, &w);
            add(u, v, w), add(v, u, w);
        }
        tot = Ma[rt] = n; //首次先找整棵树的重心
        getrt(1, 0);
        slove(rt);
        ll M = 1ll * n * n;
        ll g = __gcd(ans, M);
        printf("%lld/%lld\n", ans/g, M/g);
    	return 0;
    }

[https_www.luogu.com.cn_problem_P2634]: https://www.luogu.com.cn/problem/P2634