51nod2553 双重祖先

题意

给定两颗有根树，两颗树均有 \(n\) 个节点，且跟均为 \(1\) 号点

问有多少对 \((u,v)\) 满足在给定的两颗树中 \(u\) 均为 \(v\) 的祖先

解法

先重链剖分第一颗树，处理出剖分序后在第二颗树上 \(dfs\)

每 \(dfs\) 到一个点就把其加入第一颗树对应的剖分序的位置，可以用树状数组维护

那么假设现在 \(dfs\) 到的点为点 \(x\)，那么其祖先一定都在树状数组中，那么只要在第一颗树上查询 \(x\) 到根上 \(1\) 的个数即可。这里的 \(1\) 代表着既在第一颗树上是 \(x\) 的祖先，又在第二颗树上是 \(x\) 的祖先

代码

#include <cstdio>
#include <vector>

using namespace std;

const int MAX_N = 1e5 + 10;

void dfs_1(int x);
void dfs_2(int x, int tp);

int N;

long long ans;

int sz[MAX_N], fa[MAX_N], son[MAX_N], dep[MAX_N], dfn[MAX_N], top[MAX_N], pos[MAX_N], Id;

vector<int> e[MAX_N], g[MAX_N];

struct BIT {
    int c[MAX_N];
    void insert(int x, int v) {
        for (; x && x <= N; x += x & -x)  c[x] += v;
    }
    int query(int x, int res = 0) {
        for (; x; x -= x & -x)  res += c[x];
        return res;
    }
} tr;

int query(int x) {
    int res = 0;
    while (x) {
        res += tr.query(pos[x]) - tr.query(pos[top[x]] - 1);
        x = fa[top[x]];
    }
    return res;
}

void DFS(int x, int lst) {
    ans += query(fa[x]);
    tr.insert(pos[x], 1);
    for (int i = 0, s = g[x].size(); i < s; ++i) {
        int v = g[x][i];
        if (v ^ lst)  DFS(v, x);
    }
    tr.insert(pos[x], -1);
}

int main() {
    
    scanf("%d", &N);
    
    int u, v;
    for (int i = 1; i < N; ++i) {
        scanf("%d%d", &u, &v);
        e[u].push_back(v), e[v].push_back(u);
    }
    
    for (int i = 1; i < N; ++i) {
        scanf("%d%d", &u, &v);
        g[u].push_back(v), g[v].push_back(u);   
    }
    
    dfs_1(1);
    dfs_2(1, 1);
    
    DFS(1, 0);
    
    printf("%lld\n", ans);
    
    return 0;   
}

void dfs_1(int x) {
    sz[x] = 1;
    for (int i = 0, s = e[x].size(); i < s; ++i) {
        int v = e[x][i];
        if (v ^ fa[x]) {
            fa[v] = x, dep[v] = dep[x] + 1, dfs_1(v), sz[x] += sz[v];
            if (sz[v] > sz[son[x]])  son[x] = v;
        }
    }
}

void dfs_2(int x, int tp) {
    top[x] = tp, dfn[++Id] = x, pos[x] = Id;
    if (son[x]) {
        dfs_2(son[x], tp);
        for (int i = 0, s = e[x].size(); i < s; ++i) {
            int v = e[x][i];
            if (!pos[v])  dfs_2(v, v);
        }
    }
}