[ 实验舱 CSP/NOIP新赛制内部挑战赛4 ] 树上询问

Problem

超级加强版 题目地址

模拟赛数据 ：(n le 5000, m le 5000)

lxl 出的题，爱了爱了。

Solution

我们可把距离 (p) 为 (d) 距离的点看做一个点集，于是题目转化为告诉我们两个点集 (A,B)，求 $$sum_{a in A} sum_{b in B} dis(a,b)$$

树上路径距离 (dis(a,b) = dep[a] + dep[b] - 2*dep[lca])。于是就是求 $$sum_{a in A} sum_{b in B} dep[a] + dep[b] - 2*dep[lca]$$

把 (sum dep[a]+dep[b]) 和 (sum 2*dep[lca]) 分开来算。两者统计均可 (O(n)) 做到，时间复杂度 (O(nm))

Code

Talk is cheap.Show me the code.

#include<bits/stdc++.h>
#define int long long
using namespace std;
inline int read() {
    int x = 0, f = 1; char ch = getchar();
    while(ch<'0' || ch>'9') { if(ch=='-') f=-1; ch=getchar(); }
    while(ch>='0'&&ch<='9') { x=(x<<3)+(x<<1)+(ch^48); ch=getchar(); }
    return x * f;
}
const int N = 5007;
int n,cnt,ans;
int head[N],dep[N],sa[N],sb[N];
struct Edge {
	int next,to;
}edge[N<<1];
inline void add(int u,int v) {
	edge[++cnt] = (Edge)<%head[u], v%>;
	head[u] = cnt;
}
void Dfs1(int u,int fa) {
	dep[u] = dep[fa] + 1;
	for(int i=head[u];i;i=edge[i].next) {
		int v = edge[i].to;
		if(v != fa) {
			Dfs1(v,u);
		}
	}
}
void Bfs(int *vis,int s,int d) {
	queue<int> q; q.push(s); vis[s] = 1;
	for(int i=1;i<=d;++i) {
		int step = q.size();
		while(step--) {
			int u = q.front(); q.pop();
			for(int i=head[u];i;i=edge[i].next) {
				int v = edge[i].to;
				if(!vis[v]) {
					q.push(v); vis[v] = 1;
				}
			}
		}
	}
}
bool vis[N];
void Bfs2(int elsz,int s,int d) {
	memset(vis, 0, sizeof(vis));
	queue<int> q; q.push(s); vis[s] = 1; ans += dep[s] * elsz;
	for(int i=1;i<=d;++i) {
		int step = q.size();
		while(step--) {
			int u = q.front(); q.pop();
			for(int i=head[u];i;i=edge[i].next) {
				int v = edge[i].to;
				if(!vis[v]) {
					ans += dep[v] * elsz;
					q.push(v); vis[v] = 1;
				}
			}
		}
	}
}
void Dfs2(int u,int fa) {
	for(int i=head[u];i;i=edge[i].next) {
		int v = edge[i].to;
		if(v != fa) {
			Dfs2(v,u); sa[u] += sa[v], sb[u] += sb[v];
		}
	}
	int val = dep[u], sumA = 0;
	for(int i=head[u];i;i=edge[i].next) {
		int v = edge[i].to;
		if(v != fa) {
			ans -= val * sa[v] * (sb[u]-sb[v]);
			sumA += sa[v];
		}
	}
	ans -= val * (sa[u] - sumA) * sb[u];
}
signed main()
{
	n = read();
	for(int i=2;i<=n;++i) {
		int fa = read(); add(i,fa), add(fa,i);
	}
	Dfs1(1,0);
	int q = read();
	while(q--) {
		int p0 = read(), d0 = read(), p1 = read(), d1 = read();
		memset(sa, 0, sizeof(sa));
		memset(sb, 0, sizeof(sb));
		Bfs(sa, p0, d0); Bfs(sb, p1, d1);
		ans = 0;
		Dfs2(1,0); ans *= 2;
		Bfs2(sa[1], p1, d1); Bfs2(sb[1], p0, d0);
		printf("%lld
",ans);
	}
    return 0;
}
/*
7
1 1 2 3 5 2
1
5 1 5 0

2
*/

Summary

• 树上距离统计问题尝试用距离公式，分开计算

• 树上路径用LCA分类