Strange Memory Gym - 102832F(树上启发式合并+二进制拆分算贡献)

题目链接

题意

1. 给我们一颗 n 个节点的树，每个节点有一个权值 a [i], 现在让我们求：树上点对 i，j 满足下面的公式的个数

   

思路

1. 我们考虑有树上启发合并去暴力求解答案 O (nlogn),
2. 但是我们发现直接这样做难以维护异或和信息，这个时候我们需要维护可以考虑对每节点的编号设为 i，我们对 i 每一个二进制位统计贡献，复杂度变为 20*nlog (n)
3. 我们需要一个数组 dp [val][0/1][0~20], 第一维表示当 $ai异或a_{lca(i,j)=val==a_j}$ ，dp [val][x][k] 表示当前所求的子树中节点的值为 val，它们的第 k 个二进制位是 x 的个节点个数。

注意：两个 1e6 级别的数 异或结果可能是 2e6 级别的

代码

#include <bits/stdc++.h>

using namespace std;

const int MAXN = 1e5 + 10;

long long ans = 0;
int siz[MAXN], son[MAXN], a[MAXN];
int dp[MAXN * 12][2][20];
vector <int> E[MAXN];

void dfs1(int u, int fa)
{
    siz[u] = 1;  // 处理轻重儿子
    for (auto v : E[u])
    {
        if (v == fa)
            continue;
        dfs1(v, u);
        siz[u] += siz[v];
        if (!son[u] || siz[v] > siz[son[u]])
            son[u] = v;
    }
}

void get_ans(int u, int fa, int FA)
{
    for (int i = 0; i <= 19; ++i)
    {
        ans += (1ll << i) * dp[(a[u] ^ a[FA])][((u >> i) & 1) ^ 1][i];
    }
    for (auto v : E[u])
    {
        if (v == fa)
            continue;
        get_ans(v, u, FA);
    }
}

void add_del(int u, int fa, int op)
{
    for(int i = 0; i <= 19; ++i)
    {
        dp[a[u]][(u >> i) & 1][i] += op;
    }
    for (auto v : E[u])
    {
        if (v == fa)
            continue;
        add_del(v, u, op);
    }
}

void dfs2(int u, int fa, int keep)
{
    for (auto v : E[u])
    {
        if (v == fa || v == son[u]) // 跑轻重儿子
            continue;
        dfs2(v, u, 0);
    }
    if (son[u])
        dfs2(son[u], u, 1);
    for(int i = 0; i <= 19; ++i) // u的贡献
    {
        dp[a[u]][(u >> i) & 1][i] ++;
    }
    for (auto v : E[u])
    {
        if (v == fa || v == son[u])
            continue;
        get_ans(v, u, u); // 得到贡献
        add_del(v, u, 1); // 合并
    }
    if (!keep)
        add_del(u, fa, -1); // 暴力删除 消除贡献
}

int main()
{
    int n;
    scanf("%d", &n);
    for(int i = 1; i <= n; ++i)
    {
        scanf("%d", &a[i]);
    }
    for(int i = 0; i < n - 1; ++i)
    {
        int u, v;
        scanf("%d%d", &u, &v);
        E[u].emplace_back(v);
        E[v].emplace_back(u);
    }

    dfs1(1, -1);
    dfs2(1, -1, 0);
    printf("%lld\n", ans);
    return 0;
}