[nowcoder1103B]路径计数机

题目

传送门 to nowcoder

题目描述
有一棵 $n$ 个点的树和两个整数 $p,q$，求满足以下条件的四元组 $(a,b,c,d)$ 的个数：

• $1⩽a,b,c,d⩽n$ 。
• 点 $a$ 到点 $b$ 经过的边数为 $p$，即 $dis(a,b)=p$ 。
• 点 $c$ 到点 $d$ 经过的边数为 $q$，即 $dis(c,d)=q$ 。
• 不存在一个点，它既在点 $a$ 到点 $b$ 的路径上，又在点 $c$ 到点 $d$ 的路径上，即 $⟨a,b⟩,⟨c,d⟩$ 是点不交的。

数据范围与约定
$1\le n,p,q\le 3000$，$1\le u,v\le n$，保证给出的是一棵合法的树。

思路

首先我们要解决 怎样判断两条路径是否相交。

一般树上路径问题都得用 $lca$ 解决，此题也不例外。

考虑 $a,b$ 的 $lca$，设为 $x$，再设 $lca(c,d)=y$ 。如果 $x,y$ 不是祖先$\text{-}$后代关系，两条路径一定不交。

否则，假设 $x$ 是 $y$ 的祖先，当且仅当 $a$（或 $b$）在 $y$ 的子树内时，两条路径相交。

如图所示，满足 $a$ 在 $y$ 的子树内，于是两条路径相交。

这提示我们用 $f(x)$ 表示 $lca$ 为 $x$ 的点对中，距离为 $p$ 的数量。形式化地说，
$$f(x)=\sum _{⟨a,b⟩}[lca(a,b)=x\wedge dis(a,b)=p]$$

如此一来，所有 $lca$ 在 $a$ 到 $b$ 的路径上的点对都不可以计入答案。与之相对的，$x$ 子树内的其他点作为 $lca$ 都是可接受的。这玩意儿可以树上差分。

$y$ 是 $x$ 的祖先怎么办？发现在这样的情况下 $x$ 与其父节点的边 一定会被覆盖到。那么我们再处理覆盖一条边的就是了！

代码

#include <cstdio>
#include <iostream>
#include <vector>
using namespace std;
inline int readint(){int a; cin >> a; return a;}

struct Edge{
	int to, nxt; Edge(){}
	Edge(int T,int N):to(T),nxt(N){}
};

const int MaxN = 3001;
int n, p, q;
Edge edge[MaxN<<1]; int head[MaxN], cntEdge;

void addEdge(int from,int to){
	if(not cntEdge) for(int i=1; i<=n; ++i)
		head[i] = -1;
	edge[cntEdge] = Edge(to,head[from]);
	head[from] = cntEdge ++;
}
void init(){
	n = readint(), p = readint(), q = readint();
	for(int i=1,u,v; i<n; ++i){
		u = readint();
		v = readint();
		addEdge(u,v), addEdge(v,u);
	}
}

namespace UnionFindSet{
	int fa[MaxN];
	int findSet(int a){
		if(fa[a] != a)
			fa[a] = findSet(fa[a]);
		return fa[a];
	}
}
int lca[MaxN][MaxN], depth[MaxN];
void tarjan(int x,int pre){ // 预处理每一对lca
	UnionFindSet::fa[x] = x;
	for(int i=head[x]; ~i; i=edge[i].nxt)
		if(edge[i].to != pre){
			depth[edge[i].to] = depth[x]+1;
			tarjan(edge[i].to,x);
			UnionFindSet::fa[edge[i].to] = x;
		}
	for(int i=1; i<=n; ++i)
		if(UnionFindSet::fa[i])
			lca[x][i] = lca[i][x] = UnionFindSet::findSet(i);
}
int distance(int a,int b){ // 树上距离
	return depth[a]+depth[b]-(depth[lca[a][b]]<<1);
}
long long dfsRes[MaxN]; // 按照距离分类，计数
void dfs(int x,int deep,int pre){
	++ dfsRes[deep];
	if(deep == q) return ;
	for(int i=head[x]; ~i; i=edge[i].nxt)
		if(edge[i].to != pre)
			dfs(edge[i].to,deep+1,x);
}
long long cnt1[MaxN]; // lca为i的 
void DFS(int x,int pre){ // 树上差分，统计一条链
	if(pre >= 1) cnt1[x] += cnt1[pre];
	for(int i=head[x]; ~i; i=edge[i].nxt)
		if(edge[i].to != pre)
			DFS(edge[i].to,x);
}
long long cnt2[MaxN], tmp[MaxN]; // dis: q 
int Father[MaxN];
void solve(){
	tarjan(1,-1);
	for(int i=1; i<=n; ++i) // cnt1即为f(x)
		for(int j=1; j<=n; ++j)
			if(distance(i,j) == q)
				++ cnt1[lca[i][j]];
	long long inTotal = cnt1[1];
	for(int i=2; i<=n; ++i){
		inTotal += cnt1[i];
		for(int j=0; j<=q; ++j)
			dfsRes[j] = 0;
		for(int j=head[i]; ~j; j=edge[j].nxt)
			if(depth[edge[j].to] < depth[i]){
				Father[i] = edge[j].to; break;
			} // 找父亲大型公益活动
		dfs(i,0,Father[i]);
		for(int j=0; j<=q; ++j)
			tmp[j] = dfsRes[j], dfsRes[j] = 0;
		dfs(Father[i],0,i);
		for(int j=0; j<q; ++j)
			cnt2[i] += (tmp[j]*dfsRes[q-j-1])<<1;
		// cnt2计算，覆盖i与i父亲的边的点对
	}
	DFS(1,-1); // 树上差分
	long long ans = 0;
	for(int i=1; i<=n; ++i)
		for(int j=1; j<=n; ++j){
			if(distance(i,j) != p)
				continue;
			int x = lca[i][j];
			ans += inTotal-cnt1[i]-cnt1[j]+cnt1[Father[x]]+cnt1[x]-cnt2[x];
		}
	printf("%lld\n",ans);
}

int main(){
	init(), solve(); return 0;
}