MX 暑假集训 8.3

【初音ミク】存活千年【いよわ】【中文CC字幕】

跑大样例只跑了 \(6\) s，我还以为我复杂度没错结果是漏了一个小分支导致变成了 \(n^2\)，直接挂 \(45\) pts。

T2T3 是 SAM 和树分块，没学过不会喵。

T1: 王国

题意

给定一棵 \(n\) 个节点的树，每个点的颜色为黑白中的一种，\(q\) 次询问，每次询问给定 \(x,y,k\)，询问 \(x\) 到 \(y\) 简单路径上所有极长颜色段中有多少个长度大于 \(k\)。

\(1\le n,q\le 10^6\)。

solution

考虑每个点 \(u\) 维护深度最小的祖先 \(f_u\)，满足 \(f_u\) 到 \(u\) 的路径上仅包含一种颜色。

这个可以用并查集维护，那么对于每次询问，\(x,y\) 分别向上跳到 \(\operatorname{lca}(x,y)\) 处，维护跳的过程中经过的颜色段有几个长度大于 \(k\)，特殊处理一下 \(\operatorname{lca}(x,y)\) 的连通块，这样做是 \(O(n^2)\) 的。

考虑用树剖维护，跳的时候直接跳到重链顶端，并计算当前重链的贡献，由于最多跳 \(O(\log n)\) 次，计算贡献可以离线树状数组，于是得到了一个 \(O(n\log^2 n)\) 的做法，常数很小可以通过。

也可以不用树剖维护，实际上我们在做的事情是从 \(x\) 跳到 \(fa_{f_x}\)，那么我们可以新建一棵树，点 \(x\) 的父亲为 \(fa_{f_x}\)，边权为 \(x\) 跳到 \(fa_{f_x}\) 经过的节点数，那么在新树上跳父亲并统计过程中边权大于 \(k\) 的边的条数，并特殊处理 \(\operatorname{lca}(x,y)\) 的贡献即可，运用差分树状数组可以做到 \(O(n\log n)\)。

upd：似乎这个做法并不需要树状数组，可以做到求 \(\operatorname{lca}(x,y)\) 外线性。

Code

#include<cstdio>
#include<algorithm>
#include<vector>
#include<bitset>
using namespace std;
#define ll long long
#define qwq Ff472130
#define f(i,l,r) for (int i=l;i<=r;i++)
#define F(i,l,r) for (int i=l;i>=r;i--)
constexpr int N=1e6+10;
constexpr int inf=1e9+10;

inline void read(int &x) {
	x=0;
	char ch=getchar();
	while (ch<48) ch=getchar(); 
	while (ch>=48) x=(x<<3)+(x<<1)+(ch^48),ch=getchar();
}

int n,m;
int c[N],ans[N];
struct ques{int x,y,k,id;};

int head[N],Etot;
struct Edge{int next,to;}e[N<<1];

inline void add(int u,int v) {
	e[++Etot]={head[u],v};
	head[u]=Etot;
}

int f[N],dep[N],siz[N];
int top[N],son[N],dn[N];

inline void dfs1(int now,int fa) {
	dep[now]=dep[f[now]=fa]+(siz[now]=1);
	for (int i=head[now];i;i=e[i].next) {
		int v=e[i].to;
		if (v==fa) continue;
		dfs1(v,now);
		siz[now]+=siz[v];
		if (siz[v]>siz[son[now]]) son[now]=v;
	}
}

inline void dfs2(int now,int tp) {
	top[now]=tp;
	dn[now]=now;
	if (!son[now]) return;
	dfs2(son[now],tp);
	if (c[son[now]]==c[now]) dn[now]=dn[son[now]];
	for (int i=head[now];i;i=e[i].next) {
		int v=e[i].to;
		if (v==f[now]||v==son[now]) continue;
		dfs2(v,v);
	}
}

int Ff[N];
inline int find(int x) {
	if (Ff[x]==x) return x;
	return Ff[x]=find(Ff[x]);
}

bitset<N> did;
struct ques2{int id,k,op;};
vector<ques2> vec[N];

inline int lca(int x,int y) {
	while (top[x]!=top[y]) {
		if (dep[top[x]]<dep[top[y]]) swap(x,y);
		x=f[top[x]];
	}
	return dep[x]<dep[y]?x:y;
}

struct Node{int sum,len;};
inline Node work(int x,int y,int k,int id) {
	int fy=find(y),sum=0;
	while (1) {
		int fx=find(x),tp=top[x];
		if (fx==fy) return {sum,dep[x]-dep[y]+1};
		if (dep[tp]<dep[fx]&&dep[tp]>dep[y]&&dep[x]>dep[dn[tp]]&&c[son[x]]!=c[x]) {
			vec[x].push_back({id,k,1});
			did.set(tp);
			x=dn[tp];
		}
		else if (dep[tp]<dep[fx]&&dep[tp]<=dep[y]&&c[son[x]]!=c[x]&&dep[dn[y]]<dep[x]) {
			vec[x].push_back({id,k,1});
			did.set(tp);
			x=dn[y];
			vec[x].push_back({id,k,-1});
		}
		else {
			sum+=(dep[x]-dep[fx]+1>=k);
			x=f[fx];
		}
	}
	return {-1,-1};
}

inline int update(int x,int y,int k,int id) {
	int z=lca(x,y);
	Node k1=work(x,z,k,id),k2=work(y,z,k,id);
	return k1.sum+k2.sum+(k1.len+k2.len-1>=k);
}

struct BIT {
	int v[N];
	inline void add(int p,int k) {for(;p<=n;p+=(p&-p))v[p]+=k;}
	inline int qry(int p) {int r=0;for(;p;p-=(p&-p))r+=v[p];return r;}
}tr;

inline void solve(int x) {
	if (!did[x]) return;
	x=son[dn[x]];
	int len=0;
	vector<int> did;
	while (x) {
		len++;
		if (!son[x]||c[son[x]]!=c[x]) {
			tr.add(n-len+1,1);
			did.push_back(n-len+1);
			len=0;
		}
		for (ques2 p:vec[x]) ans[p.id]+=tr.qry(n-p.k+1)*p.op;
		x=son[x];
	}
	for (int k:did) tr.add(k,-1);
}

int main() {
	freopen("kingdom.in","r",stdin);
	freopen("kingdom.out","w",stdout);
	read(n);
	f(i,1,n) read(c[i]);
	f(i,2,n) {
		int u,v;read(u);read(v);
		add(u,v);add(v,u);
	}
	dfs1(1,0);dfs2(1,1);Ff[1]=1;
	f(i,2,n) {
		if (c[f[i]]==c[i]) Ff[i]=f[i];
		else Ff[i]=i;
	}
	read(m);
	f(i,1,m) {
		int x,y,k;read(x);read(y);read(k);
		ans[i]=update(x,y,k,i);
	}
	f(i,1,n) if (top[i]==i) solve(i);
	f(i,1,m) printf("%d\n",ans[i]);
	return 0;
}