51nod2553 双重祖先

题意

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

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

解法

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

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

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

代码

#include 
#include 

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 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);
		}
	}
}