F. Strange Memory

题目链接

• 题意：给定一个结点个数为 $n$ 的根为 $1$ 的有根树，每个结点有一个权值 $a[i]$。我们需要找到所有满足 $a[u]\oplus a[v]==a[lca]$ 的点对 $u,v$ ，统计所有这些点对 $u\oplus v$ 的和。

思路：

因为结点权值不能为 $0$，所以一定不会出现 $a[u]\oplus a[v]==a[u]$ 的情况，满足条件的点对一定是在某 $lca$ 结点的两个不同的子树上。所以我们很容易可以想到枚举 $lca$，将该 $lca$ 结点下的子树分成两个集合 A 和 B，保证集合 A 和 B 中的结点不重合且 A 中的结点和 B 中的结点一定在 $lca$ 结点下的两棵不同的子树上。这样的话我们再枚举集合 A 中的结点 u，找到集合 B 中所有权值为 a[u] ^ a[lca] 的结点并统计答案是不是就可以了？

如果暴力枚举是 n * n 的复杂度，所以想想怎么优化呢？

想到异或的一个性质：$a\oplus b$ 的结果等于将 $a$ 和 $b$ 拆分成二进制按位异或后结果的和。
比如1101 ^ 1001 == 1 000 ^ 1 000 + 1 00 ^ 0 00 + 0 0 ^ 0 0 + 1 ^ 1 。
这样的话我们就可以记录下权值为 $val$ 的所有结点 $u$ ：直接将 $u$ 拆分成二进制记录 $val$ 下每一位的个数即$num[val][i][(u>>i)\&1]$ 。

这样的话，我们可以统计每棵子树的贡献，于是想到用 $dsuontree$。将枚举 $u$ 优化到 $log$ 的复杂度。针不戳！我们处理集合问题的时候，对于子树 $lca$，开始集合 $B$ 为空，每遍历一棵 $lca$ 的子树统计贡献，就将这棵子树的结点放到集合 $B$ 里。

#include <iostream>
#include <cstdio>
#include <vector>
#include <cstring>

typedef long long ll ;
using namespace std;
const int maxN = (1 << 20); //1e6 < 2 ^ 20, 1e6内的值异或最大为2^20-1(二进制20位都为1)

int read()
{
    int x = 0, f = 1; char c = getchar();
    while(c < '0' || c > '9') { if(c == '-') f = -f; c = getchar(); }
    while(c >= '0' && c <= '9') { x = x * 10 + c - '0'; c = getchar(); }
    return x * f;
}

struct EDGE{
    int adj, to;
    EDGE() {}
    EDGE(int a, int b) : adj(a), to(b) {}
}edge[maxN << 1];
int head[maxN], cnt;
void init(int n) {
    for(int i = 0; i <= n; ++ i ) {
        head[i] = -1;
    }
    cnt = 0;
}
void add_edge(int u, int v) {
    edge[cnt] = EDGE(head[u], v);
    head[u] = cnt ++ ;
}

int n, a[maxN];
int siz[maxN], son[maxN], nowSon;
void dfs(int u, int fa) {
    siz[u] = 1;
    for(int i = head[u]; ~i; i = edge[i].adj) {
        int v = edge[i].to;
        if(v == fa) continue;
        dfs(v, u);
        siz[u] += siz[v];
        if(siz[son[u]] < siz[v]) son[u] = v;
    }
}
int num[maxN][20][2];
ll ans;
void add(int u, int val) { //更新贡献a[u]
    for(int i = 0; i < 20; ++ i ) {
        num[a[u]][i][(u >> i) & 1] += val;
    }
}

void del(int u, int fa) { //删除轻子树
    add(u, -1);
    for(int i = head[u]; ~i; i = edge[i].adj) {
        int v = edge[i].to;
        if(v == fa) continue;
        del(v, u);
    }
}

ll query(int u, int val) { //查询结点u能够匹配到的答案：a[?] ^ a[u] == val的所有? ^ u的和
    ll res = 0;
    for(int i = 0; i < 20; ++ i ) {
        int bit = !((u >> i) & 1);
        res += (ll)num[val][i][bit] * ll(1 << i);
    }
    return res;
}

vector<int>vt;
void cal(int u, int fa, int lca) {
    vt.emplace_back(u);
    ans += query(u, a[u] ^ a[lca]);
    for(int i = head[u]; ~i ; i = edge[i].adj) {
        int v = edge[i].to;
        if(v == fa || v == nowSon) continue;
        cal(v, u, lca);
    }
}

void dsu(int u, int fa, bool keep) {
    for(int i = head[u]; ~i; i = edge[i].adj) {
        int v = edge[i].to;
        if(v == fa || v == son[u]) continue;
        dsu(v, u, false);
    }
    if(son[u]) {
        dsu(son[u], u, true);
        nowSon = son[u];
    }

    for(int i = head[u]; ~i; i = edge[i].adj) {
        int v = edge[i].to;
        if(v == fa || v == nowSon) continue;
        cal(v, u, u); //暴力统计轻儿子的贡献， 知道lca，知道u，找集合中现有合法的v
        //统计完一棵子树之后再更新集合中的结点
        for(auto j : vt) {
            add(j, 1);
        }
        vector<int>().swap(vt);
    }

    add(u, 1); //将当前根结点u加入集合中
    nowSon = 0;

    if(!keep) {
        del(u, fa);
    }
}

int main()
{
    n = read();
    init(n);
    for(int i = 1; i <= n; ++ i ) {
        a[i] = read();
    }
    for(int i = 1; i < n; ++ i ) {
        int u = read(), v = read();
        add_edge(u, v);
        add_edge(v, u);
    }
    dfs(1, 0);
    dsu(1, 0, true);
    cout << ans << endl;
    return 0;
}