[NOI2021] 轻重边题解

题目传送门

一眼数据结构

考虑树上有什么数据结构支持 \(x\) 到 \(y\) 节点的修改和查询，那就是：树链剖分。

那么这道树链剖分的题有个 \(trick\)：边点转换&染色法，对于每次修改，考虑将修改路径上的点全部染上一种独一无二的颜色，而查询的时候，就查询路径上相邻的相同的颜色节点个数即可。

正确性易证：因为每次修改，该点周围的边都变成轻边除了那条路径上的边，因为那条路径上的点必然与该点颜色一样，而不是路径上的点必然与该点颜色不一样，所以可以转换成求路径上相邻的相同的颜色节点个数。

考虑怎么求：对于线段树上的点，开一个结构体，记 \(lc:\) 该区间左端点颜色，\(rc:\) 该区间右端点颜色，\(gs:\) 该区间相邻相同颜色节点个数，\(tag:\) 懒标记（要区间修改）

一些小细节：

\(1\)，初始化时把所有数组都初始化了，防止出现神奇错误（比如我为此调了半小时），不差这点常数。

\(2\)，要特别注意 \(top_x=top_y\) 的情况，大部分错误都出现在这里。

上代码：

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const ll N=1e5+50;
ll T,n,m,u,v;
vector<ll> e[N];
ll dep[N],fa[N],siz[N],son[N];
void dfs1(ll wz,ll last)
{
	dep[wz]=dep[last]+1;
	fa[wz]=last;
	siz[wz]=1;
	ll bigson=0;
	for(ll i=0;i<e[wz].size();i++)
	{
		ll j=e[wz][i];
		if(j==last) continue;
		dfs1(j,wz);
		siz[wz]+=siz[j];
		if(siz[j]>bigson) bigson=siz[j],son[wz]=j;
	}
}
ll top[N],id[N],cnt;
void dfs2(ll wz,ll topf)
{
	top[wz]=topf;
	id[wz]=++cnt;
	if(!son[wz]) return ;
	dfs2(son[wz],topf);
	for(ll i=0;i<e[wz].size();i++)
	{
		ll j=e[wz][i];
		if(j==fa[wz]||j==son[wz]) continue;
		dfs2(j,j);
	}
}
struct jgt
{
	ll lc,rc,gs,tag;
}tr[N*8];
ll idx;
void pushup(jgt &wz,jgt l,jgt r)
{
	wz.gs=l.gs+r.gs+(l.rc==r.lc);
	wz.rc=r.rc;
	wz.lc=l.lc;
}
void BT(ll wz,ll l,ll r)
{
	if(l==r)
	{
		tr[wz].lc=l;
		tr[wz].rc=l;
		tr[wz].gs=0;
		tr[wz].tag=0;
		return ;
	}
	ll mid=(l+r)/2;
	BT(wz*2,l,mid);
	BT(wz*2+1,mid+1,r);
	pushup(tr[wz],tr[wz*2],tr[wz*2+1]);
}
void pushdown(ll wz,ll l,ll r)
{
	ll mid=(l+r)/2;
	tr[wz*2].gs=mid-l;
	tr[wz*2].rc=tr[wz].tag;
	tr[wz*2].lc=tr[wz].tag;
	tr[wz*2].tag=tr[wz].tag;
	tr[wz*2+1].gs=r-mid-1; 
	tr[wz*2+1].rc=tr[wz].tag;
	tr[wz*2+1].lc=tr[wz].tag;
	tr[wz*2+1].tag=tr[wz].tag;
	tr[wz].tag=0;
}
void gai(ll wz,ll l,ll r,ll le,ll ri,ll co)
{
	if(le<=l&&ri>=r)
	{
		tr[wz].gs=r-l;
		tr[wz].rc=co;
		tr[wz].lc=co;
		tr[wz].tag=co;
		return ;
	}
	if(tr[wz].tag&&l!=r) pushdown(wz,l,r);
	ll mid=(l+r)/2;
	if(le<=mid) gai(wz*2,l,mid,le,ri,co);
	if(ri>mid) gai(wz*2+1,mid+1,r,le,ri,co);
	pushup(tr[wz],tr[wz*2],tr[wz*2+1]);
}
jgt query(ll wz,ll l,ll r,ll le,ll ri)
{
	if(le<=l&&ri>=r) return tr[wz];
	if(tr[wz].tag&&l!=r) pushdown(wz,l,r);
	ll mid=(l+r)/2;
	if(le<=mid&&ri>mid)
	{
		jgt t1=query(wz*2,l,mid,le,ri);
		jgt t2=query(wz*2+1,mid+1,r,le,ri);
		jgt t3;
		pushup(t3,t1,t2);
		return t3;
	} 
	if(le<=mid) return query(wz*2,l,mid,le,ri);
	if(ri>mid) return query(wz*2+1,mid+1,r,le,ri);
}
void upd(ll x,ll y,ll co)
{
	while(top[x]!=top[y])
	{
		if(dep[top[x]]<dep[top[y]]) swap(x,y);
		gai(1,1,n,id[top[x]],id[x],co);
		x=fa[top[x]];
	}
	if(dep[x]>dep[y]) swap(x,y);
	gai(1,1,n,id[x],id[y],co);
}
ll updQ(ll x,ll y)
{
	ll ans1=0,ans2=0,co1=0,co2=0,w=0;
	while(top[x]!=top[y])
	{
		if(dep[top[x]]<dep[top[y]]) swap(x,y),w^=1;
		jgt tzy=query(1,1,n,id[top[x]],id[x]);
		if(w==1)
		{
			ans2=ans2+tzy.gs+(tzy.rc==co2);
			co2=tzy.lc;
		}
		else
		{
			ans1=ans1+tzy.gs+(tzy.rc==co1);
			co1=tzy.lc;
		}
		x=fa[top[x]];
	}
	if(dep[x]>dep[y]) swap(x,y),w^=1;
	jgt tzy=query(1,1,n,id[x],id[y]);
	if(w==0)
		return (ans1+ans2+tzy.gs+(tzy.lc==co1)+(tzy.rc==co2));
	else
		return (ans1+ans2+tzy.gs+(tzy.lc==co2)+(tzy.rc==co1));
}
void Init()
{
	cnt=0,idx=0;
	memset(son,0,sizeof son);
	memset(fa,0,sizeof fa);
	memset(siz,0,sizeof siz);
	memset(dep,0,sizeof dep);
	memset(top,0,sizeof top);
	memset(id,0,sizeof id);
	memset(tr,0,sizeof tr);
	for(ll i=1;i<=n;i++) e[i].clear(); 
}
int main()
{
	scanf("%lld",&T);
	while(T--)
	{
		scanf("%lld %lld",&n,&m);
		for(ll i=1;i<n;i++)
		{
			scanf("%lld %lld",&u,&v);
			e[u].push_back(v);
			e[v].push_back(u);
		}
		dfs1(1,0);
		dfs2(1,1);
		BT(1,1,n);
		idx=n;
		ll opt,x,y;
		for(ll i=1;i<=m;i++)
		{
			scanf("%lld %lld %lld",&opt,&x,&y);
			if(opt==1) upd(x,y,++idx);
			if(opt==2) printf("%lld\n",updQ(x,y));
		}
		Init();
	}
	return 0;
}