牛客提高集训营6 C 树(树链剖分)

题目链接

为了纪(zhuang)念(bi)写完这个树剖单独写一篇。感觉还好，也就6k嘛。
完整比赛题解：https://www.cnblogs.com/SovietPower/p/9826829.html。

肯定要用点来表示边的颜色，然后树剖。
对于操作2，我们可以拆成：

1. 将\(u\to v\)路径上的点，所有连向子节点的边染成\(col2\)；
2. 将\(u\to v\)路径上的点，所有连向父节点的边染成\(col1\)；
3. 将\(LCA(u,v)\)连向父节点的边染成\(col2\)。

那么本题的操作实际有四种：链修改、链查询、修改一条链上所有连向子节点的边、换根+子树修改。
除了链修改指向子节点的边，其它都是树剖模板（换根见BZOJ3083）。

因为是改所有指向子节点的边，所以也可以树剖维护。用\(to\_son[x]\)表示\(x\)指向儿子的轻边的颜色，用另一棵线段树维护（注意只是轻边，重边要单独在第一棵树上改）。
除此之外，在第一棵线段树上维护每个点\(x\)到\(fa[x]\)的边的颜色。

这样查询时，对于重链可以直接在第一棵线段树上查；对于轻边（\(top[x]\to fa[top[x]]\)）的颜色，需要区分是第一棵树上\(top[x]\)的值还是第二棵树上\(fa[top[x]]\)的值。修改时再带一个时间优先级即可。
要修改的链是一样的，且\(to\_son[x]\)只会单点查（切换重链时）（也是单点修改），所以只要一棵线段树同时维护两种标记就行了。
复杂度\(O(m\log^2n)\)。
另外单点查完全可以用非递归。
边界（\(dfn\pm 1\)）问题要注意。（会不会还有锅啊。。）

//358ms	20428KB
#include <cstdio>
#include <cctype>
#include <algorithm>
#define fir first
#define sec second
#define mp std::make_pair
#define pr std::pair<int,int>//time val
//#define gc() getchar()
#define MAXIN 300000
#define gc() (SS==TT&&(TT=(SS=IN)+fread(IN,1,MAXIN,stdin),SS==TT)?EOF:*SS++)
typedef long long LL;
const int N=1e5+5;

int n,Enum,H[N],nxt[N<<1],to[N<<1],fa[N],dep[N],sz[N],son[N],top[N],dfn[N],Index,ref[N],out[N];
char IN[MAXIN],*SS=IN,*TT=IN;
struct Segment_Tree
{
	#define ls rt<<1
	#define rs rt<<1|1
	#define lson l,m,ls
	#define rson m+1,r,rs
	#define S N<<2
	int cnt[S][3],sz[S];
	pr tag1[S],tag2[S];
	#undef S

	#define Update(rt) cnt[rt][0]=cnt[ls][0]+cnt[rs][0], cnt[rt][1]=cnt[ls][1]+cnt[rs][1], cnt[rt][2]=cnt[ls][2]+cnt[rs][2]
	#define Cov1(rt,v) cnt[rt][0]=0, cnt[rt][1]=0, cnt[rt][2]=0, cnt[rt][v.sec]=sz[rt], tag1[rt]=v
	#define Cov2(rt,v) tag2[rt]=v
	inline void PushDown(int rt)
	{
		int l=ls,r=rs;
		if(tag1[rt].fir) Cov1(l,tag1[rt]), Cov1(r,tag1[rt]), tag1[rt].fir=0;
		if(tag2[rt].fir) Cov2(l,tag2[rt]), Cov2(r,tag2[rt]), tag2[rt].fir=0;
	}
	void Build(int l,int r,int rt)
	{
		cnt[rt][0]=sz[rt]=r-l+1;
		if(l!=r)
		{
			int m=l+r>>1;
			Build(lson), Build(rson);
		}
	}
	void ModifyI(int l,int r,int rt,int L,int R,pr v)//为了常数 粘了仨Modify→_→
	{
		if(L<=l && r<=R) {Cov1(rt,v); return;}
		PushDown(rt);
		int m=l+r>>1;
		if(L<=m) ModifyI(lson,L,R,v);
		if(m<R) ModifyI(rson,L,R,v);
		Update(rt);
	}
	void Modify1(int l,int r,int rt,int p,pr v)
	{
		if(l==r) {Cov1(rt,v); return;}
		PushDown(rt);
		int m=l+r>>1;
		p<=m ? Modify1(lson,p,v) : Modify1(rson,p,v);
		Update(rt);
	}
	void Modify2(int l,int r,int rt,int L,int R,pr v1,pr v2)
	{
		if(L<=l && r<=R) {Cov1(rt,v1), Cov2(rt,v2); return;}
		PushDown(rt);
		int m=l+r>>1;
		if(L<=m) Modify2(lson,L,R,v1,v2);
		if(m<R) Modify2(rson,L,R,v1,v2);
		Update(rt);
	}
	pr QueryP(int l,int r,int rt,int p,int id)
	{
		int m;
		while(l!=r)
		{
			PushDown(rt);
			m=l+r>>1, p<=m?(r=m,rt=ls):(l=m+1,rt=rs);
		}
		return id==1?tag1[rt]:tag2[rt];
	}
	int QueryI(int l,int r,int rt,int L,int R,int c)
	{
		if(L<=l && r<=R) return cnt[rt][c];
		PushDown(rt);
		int m=l+r>>1;
		if(L<=m)
			if(m<R) return QueryI(lson,L,R,c)+QueryI(rson,L,R,c);
			else return QueryI(lson,L,R,c);
		return QueryI(rson,L,R,c);
	}
}T;

inline int read()
{
	int now=0;register char c=gc();
	for(;!isdigit(c);c=gc());
	for(;isdigit(c);now=now*10+c-'0',c=gc());
	return now;
}
inline void AE(int u,int v)
{
	to[++Enum]=v, nxt[Enum]=H[u], H[u]=Enum;
	to[++Enum]=u, nxt[Enum]=H[v], H[v]=Enum;
}
inline int Jump(int u,int lca)
{
	int las=u;
	while(top[u]!=top[lca]) u=fa[las=top[u]];
	return u==lca?las:ref[dfn[lca]+1];
}
void DFS1(int x)
{
	int mx=0; sz[x]=1;
	for(int i=H[x],v; i; i=nxt[i])
		if((v=to[i])!=fa[x])
		{
			fa[v]=x, dep[v]=dep[x]+1, DFS1(v), sz[x]+=sz[v];
			if(sz[v]>mx) mx=sz[v], son[x]=v;
		}
}
void DFS2(int x,int tp)
{
	top[x]=tp, dfn[x]=++Index, ref[Index]=x;
	if(son[x])
	{
		DFS2(son[x],tp);
		for(int i=H[x],v; i; i=nxt[i])
			if((v=to[i])!=fa[x]&&v!=son[x]) DFS2(v,v);
	}
	out[x]=Index;
}
int Query(int u,int v,int c)
{
	int ans=0,tp;
	while(top[u]!=top[v])
	{
		if(dep[top[u]]<dep[top[v]]) std::swap(u,v);
		tp=top[u];
		if(dfn[tp]<dfn[u]) ans+=T.QueryI(1,n,1,dfn[tp]+1,dfn[u],c);
		pr t1=T.QueryP(1,n,1,dfn[tp],1),t2=T.QueryP(1,n,1,dfn[fa[tp]],2);
		ans+=t1.fir>t2.fir?(t1.sec==c):(t2.sec==c);
		u=fa[tp];
	}
	if(dep[u]<dep[v]) std::swap(u,v);
	if(dfn[v]<dfn[u]) ans+=T.QueryI(1,n,1,dfn[v]+1,dfn[u],c);
	return ans;
}
void Modify(int u,int v,int c1,int c2,int t)
{
	if(dfn[u]<n) T.Modify1(1,n,1,dfn[u]+1,mp(t,c2));//dfn[son[u]]
	if(dfn[v]<n) T.Modify1(1,n,1,dfn[v]+1,mp(t,c2));//这个修改就不会有冲突了 
	while(top[u]!=top[v])
	{
		if(dep[top[u]]<dep[top[v]]) std::swap(u,v);
		T.Modify2(1,n,1,dfn[top[u]],dfn[u],mp(t+1,c1),mp(t,c2));
		u=fa[top[u]];
		T.Modify1(1,n,1,dfn[u]+1,mp(t+1,c2));
	}
	if(dep[u]<dep[v]) std::swap(u,v);
	T.Modify2(1,n,1,dfn[v],dfn[u],mp(t+1,c1),mp(t,c2));
	T.Modify1(1,n,1,dfn[v],mp(t+1,c2));
}

int main()
{
	n=read();
	for(int i=1; i<n; ++i) AE(read(),read());
	DFS1(1), DFS2(1,1), T.Build(1,n,1);
	for(int opt,x,y,c,c2,rt,Q=read(),i=1; i<=Q; ++i)
		switch(opt=read())
		{
			case 1: x=read(),y=read(),printf("%d\n",Query(x,y,read())); break;
			case 2: x=read(),y=read(),c=read(),c2=read(),Modify(x,y,c,c2,i<<1); break;
			case 3:
			{
				rt=read(),x=read(),c=read();
				if(x==rt) T.ModifyI(1,n,1,1,n,mp(i<<1,c));
				else if(dfn[x]<=dfn[rt]&&dfn[rt]<=out[x])
				{
					y=Jump(rt,x), T.ModifyI(1,n,1,1,dfn[y]-1,mp(i<<1,c));
					if(out[y]<n) T.ModifyI(1,n,1,out[y]+1,n,mp(i<<1,c));
				}
				else if(dfn[x]<out[x]) T.ModifyI(1,n,1,dfn[x]+1,out[x],mp(i<<1,c));
				break;
			}
		}

	return 0;
}