Data structure [树链剖分] [莫队] [分块]

链接

2020牛客暑期多校训练营（第六场）
https://ac.nowcoder.com/acm/contest/5671/D

题意

给定一个大小为 $n$ 的树， $m$ 次询问，查询 $x$ 的子树内编号在 $[l,r]$ 内 $lca=x$ 的点对数。$n,m\le 200000$

分析

问题可以转化为，查询 $x$ 的子树内在 $[l,r]$ 内的点对数，减去 $x$ 的每个儿子的子树在 $[l,r]$ 内的点对数。而求点对数就等价于求点个数。

对树进行树链剖分，一个询问就转换为求 $x$ 子树内 $[l,r]$ 的个数，以及重儿子和各个轻儿子的子树内 $[l,r]$ 的个数。

原始思路

计算 $x$ 点和 $x$ 的重儿子的贡献可以用线段树合并：树上每一个节点上开一棵线段树，每个点合并所有儿子的线段树，在线段树上查询，复杂度 $O(n\ast log(n))$。

用莫队计算每个询问的轻儿子的贡献，用莫队维护编号区间。
每当莫队的左右端点移动时，会添加或删除一个编号的点，这个点到根的路径上最多只会经过 $log(n)$ 条轻边，因此这个点只会影响 $log(n)$ 个点上的询问。
我们只需在每条重链的 $top$ 上记录子树内的个数，在 $top$ 的父亲上记录轻边的总贡献。

这样整体的复杂度为 $O(n\sqrt{n}\ast log(n))$。

考虑优化

可以发现，整体的复杂度主要集中在莫队端点移动时的跳链上，我们考虑减少跳链的次数。
对于一个点，将子树大小前 $O(\sqrt{n})$ 大的儿子都当作重儿子，那么轻儿子就是除了这些重儿子以外的所有儿子。
这样每次经过一条轻边，子树的大小会增大 $O(\sqrt{n})$ 倍，于是只会经过 $O(1)$ 条轻边，因此计算轻儿子贡献的复杂度就变为 $O(n\sqrt{n})$。

再计算重儿子的贡献，问题就变成查询 $O(n\sqrt{n})$ 次在某个子树内，编号在某个区间内的元素个数。
由于只有 $n$ 个点，我们可以用一个 $O(\sqrt{n})$ 修改—— $O(1)$ 查询的分块来维护编号在某个区间内的元素个数。
具体做法是：在dfs时，容器中存入所有遍历过的点。

1. 进入一个点时，将这个点上所有询问的贡献减去在容器中查询的数量。
2. 将当前的点加入到容器中，然后继续遍历子树。
3. 离开一个点时，将这个点上所有询问的贡献加上在容器中查询的数量。
   最终，总时间复杂度为 $O(n\sqrt{n})$，总空间复杂度为 $O(n)$。

代码

#include"bits/stdc++.h"
using namespace std;
typedef long long ll;
typedef double db;
template<class T>inline void MAX(T &x,T y){if(y>x)x=y;}
template<class T>inline void MIN(T &x,T y){if(y<x)x=y;}
template<class T>inline void rd(T &x){
	x=0;char o,f=1;
	while(o=getchar(),o<48)if(o==45)f=-f;
	do x=(x<<3)+(x<<1)+(o^48);
	while(o=getchar(),o>47);
	x*=f;
}
const int M=2e5+5;
const int S=10;
int n,m,rt,sz[M],fa[M];
vector<int>edge[M];
struct node{
	int l,r,id;
};
vector<node>vi[M];
void dfs(int x,int f){
	fa[x]=f;
	sz[x]=1;
	for(int i=0;i<edge[x].size();i++){
		int y=edge[x][i];
		if(y==f)continue;
		dfs(y,x);
		sz[x]+=sz[y];
	}
	sort(edge[x].begin(),edge[x].end(),[](int x,int y){return sz[x]>sz[y];});
}
struct BLOCK{
	static const int S=400;
	static const int K=M/S+5;
	int big[K],small[M],bel[M],L[K],R[K];
	BLOCK(){
		for(int i=1;i<M;i++)bel[i]=(i-1)/S+1;
		for(int i=1;i<K;i++)L[i]=1+(i-1)*S,R[i]=i*S;
	}
	void insert(int x){
		for(int i=1;i<=bel[x];i++)big[i]++;
		for(int i=L[bel[x]];i<=x;i++)small[i]++;
	}
	int query(int l,int r){
		return big[bel[l]+1]+small[l]-big[bel[r+1]+1]-small[r+1];
	}
}block;
struct MODUI{
	int res,cnt[M];
	void clear(int tot){res=0;for(int i=1;i<=tot;i++)cnt[i]=0;}
	void insert(int x){res+=cnt[x]++;}
	void erase(int x){res-=--cnt[x];}
	inline int query(){return res;}
}modui;
int now1,now2,cnt1[M],cnt2[M],id[M],reid[M],dfsid;
ll ans[M];
pair<int,int>color[M];
void chain(int x,int top){
	reid[id[x]=++dfsid]=x;
	int st1=now1,st2=now2;
	for(int i=0;i<vi[x].size();i++)
		cnt1[now1++]=-block.query(vi[x][i].l,vi[x][i].r);
	if(top!=x)for(int i=0,len=vi[fa[x]].size();i<len;i++)
		cnt2[now2++]=-block.query(vi[fa[x]][i].l,vi[fa[x]][i].r);
	block.insert(x);
	for(int i=0;i<edge[x].size();i++){
		int y=edge[x][i];
		if(y==fa[x])continue;
		if(i<S)chain(y,top);
		else chain(y,y);
	}
	for(int i=0;i<vi[x].size();i++){
		cnt1[st1+i]+=block.query(vi[x][i].l,vi[x][i].r);
		ans[vi[x][i].id]+=1ll*cnt1[st1+i]*(cnt1[st1+i]-1)/2;
	}
	if(top!=x)for(int i=0,len=vi[fa[x]].size();i<len;i++){
		cnt2[st2+i]+=block.query(vi[fa[x]][i].l,vi[fa[x]][i].r);
		ans[vi[fa[x]][i].id]-=1ll*cnt2[st2+i]*(cnt2[st2+i]-1)/2;
	}
	now1=st1,now2=st2;
	if(edge[x].size()>=S){
		int all=0,tot=0;
		for(int i=S;i<edge[x].size();i++){
			int y=edge[x][i];
			if(y==fa[x])continue;
			tot++;
			for(int j=id[y];j<id[y]+sz[y];j++)color[++all]=make_pair(reid[j],tot);
		}
		sort(color+1,color+all+1);
		for(int i=0;i<vi[x].size();i++){
			vi[x][i].l=lower_bound(color+1,color+all+1,make_pair(vi[x][i].l,0))-color;
			vi[x][i].r=upper_bound(color+1,color+all+1,make_pair(vi[x][i].r,233333))-color-1;
		}
		int S=sqrt(all)+1;
		sort(vi[x].begin(),vi[x].end(),[&](node &a,node &b){
			if(a.l/S!=b.l/S)return a.l/S<b.l/S;
			return a.r<b.r;
		});
		int L=1,R=0;
		modui.clear(tot);
		for(int i=0;i<vi[x].size();i++){
			int l=vi[x][i].l,r=vi[x][i].r,id=vi[x][i].id;
			if(l>r)continue;
			while(R<r)modui.insert(color[++R].second);
			while(L>l)modui.insert(color[--L].second);
			while(R>r)modui.erase(color[R--].second);
			while(L<l)modui.erase(color[L++].second);
			ans[id]-=modui.query();
		}
	}
}
int main(){
#ifndef ONLINE_JUDGE
	freopen("jiedai.in","r",stdin);
//	freopen("jiedai.out","w",stdout);
#endif
	rd(n),rd(m),rd(rt);
	for(int i=1;i<n;i++){
		int a,b;
		rd(a),rd(b);
		edge[a].push_back(b);
		edge[b].push_back(a);
	}
	for(int i=1;i<=m;i++){
		int l,r,x;
		rd(l),rd(r),rd(x);
		vi[x].push_back((node){l,r,i});
	}
	dfs(rt,0);
	chain(rt,rt);
	for(int i=1;i<=m;i++)printf("%lld\n",ans[i]);
	return 0;
}