yww 与树上的回文串 (点分治+AC自动机+分块决策)

yww 与树上的回文串

官方题解不是人啊，半天看不懂

首先是前置知识 Border定理，在这道题的应用就是把一个串的回文的前缀，按照长度分成\(O(\log |S|)\)个等差数列

分析

路径问题很常见的可以考虑点分治，问题在于如何合并两个子树里的串回文的个数

如果符合条件的串只有在一边，可以直接\(dfs\),\(\text{hash}\)判断

设符合条件的串被当前点分治的根分成了两部分\(S_x\),\(S_y\)，且\(|S_x|\ge |S_y|\)

其中\(|S_x|=|S_y|\)的部分可以特判一下

可以认为，合法的情况就是\(S_x\)长度为\(|S_x|-|S_y|\)的前缀自身回文，剩余的部分与\(S_y\)相同

那么我们枚举每个\(S_x\)，建立AC自动机，用\(\text{hash}\)同步维护每个前缀是否回文，插入等差数列中，得到\(log|S|\)个等差数列

接下来对于公差\(d\)我们分段决策

每一段等差数列可以表示为三元组\((l,r,d)\)，表示剩余的部分长度\(x\)的集合是\(x\in[l,r],x \equiv l(\mod d)\)

而剩余的部分是当前串在\(fail\)树上的祖先，所以问题可以转化为一个多条件的树上前缀查询

设常数\(S\)为决策边界

如果\(d\leq S\)，我们把每一个\(d\)的答案都统计下来，建立\(S^2\)棵虚空线段树，就能查询区间\([l,r]\)的答案了

树上前缀查询可以离线询问，记录\(\text{trie}\)树上每个点作为\(|S_y|\)结尾的个数，枚举\(S\)种不同的模数更新，遍历\(fail\)树就能得到

设当前点分治包含点个数为\(n\)，这一部分复杂度是\(O(nS\log n)\)

如果\(d>S\)，直接用一个长度为\(n\)的数组记录树上前缀中长度为\(i\)的答案，枚举集合\((l,r,d)\)即可

这一部分复杂度是\(O(\frac{n^2}{S})\)

实测\(S=4,5\)时效果奇好。。

const int N=5e4+10,SN=240;
const ll K1=234234,K2=133123,P1=1e9+13,P2=19260817;

int n;
vector <pair <int,int> > G[N];
vector <int> E[N];
ll Pow1[N],Pow2[N];
int trie[N][2],End[N],cnt,fail[N],dep[N]; // End数组记录当前节点作为Sy结尾的个数
struct Node{ int l,r,d; };
vector <Node> Q[N];
void Clear(){
	rep(i,0,cnt) fail[i]=trie[i][0]=trie[i][1]=End[i]=0,Q[i].clear(),E[i].clear();
	cnt=0;
}
void Build(){
	static queue <int> que;
	rep(i,0,1) if(trie[0][i]) que.push(trie[0][i]),dep[trie[0][i]]=1;
	while(!que.empty()) {
		int u=que.front(); que.pop();
		E[fail[u]].pb(u);
		rep(i,0,1) {
			int &v=trie[u][i];
			if(!v) v=trie[fail[u]][i];
			else dep[v]=dep[u]+1,fail[v]=trie[fail[u]][i],que.push(v);
		}
	}
}

int vis[N],mi,Rt,sz[N];
void FindRt(int n,int u,int f) {
	int ma=n-sz[u];
	for(auto t:G[u]) if(!vis[t.first] && t.first!=f) FindRt(n,t.first,u),cmax(ma,sz[t.first]);
	if(mi>ma) mi=ma,Rt=u;
}

int Maxlen,M;
ll ans;
int a[30][N],ac; //Palindrome Prefix 

void pre_dfs(int u,int f,int d) {
	cmax(Maxlen,d+2),sz[u]=1;
	for(auto t:G[u]) if(!vis[t.first] && t.first!=f){
		int v=t.first;
		pre_dfs(v,u,d+1);
		sz[u]+=sz[v];
	}
}

void dfs(int u,int f,int now,int d,ll h1,ll h2,ll rh1,ll rh2,int k) {
	if(now) End[now]++;
	if(now && ac) 
        rep(i,1,ac) 
        	Q[now].pb((Node){d-a[i][a[i][0]],d-a[i][1],(a[i][0]>1?a[i][2]-a[i][1]:1)});
	if(now && h1==rh1 && h2==rh2) {
		ans+=k;
		if(a[ac][0]==1 || a[ac][2]-a[ac][1]==d-a[ac][a[ac][0]]) a[ac][++a[ac][0]]=d;
		else ++ac,a[ac][a[ac][0]=1]=d;
	} //同步维护等差数列
	for(auto t:G[u]) if(!vis[t.first] && t.first!=f){
		int v=t.first,w=t.second;
		if(!trie[now][w]) trie[now][w]=++cnt;
		dfs(v,u,trie[now][w],d+1,(h1*K1+w)%P1,(h2*K2+w)%P2,(rh1+Pow1[d]*w)%P1,(rh2+Pow2[d]*w)%P2,k);
	}
	if(now && h1==rh1 && h2==rh2) if(--a[ac][0]==0) --ac; // 回溯等差数列
}

const int K=30;
struct I_Hate_It{
	int ls[N*K],rs[N*K],s[N*K],cnt;
	void Upd(int &p,int l,int r,int x,int y) {
		if(!p) p=++cnt,ls[cnt]=rs[cnt]=s[cnt]=0;
		s[p]+=y;
		if(l==r) return;
		int mid=(l+r)>>1;
		x<=mid?Upd(ls[p],l,mid,x,y):Upd(rs[p],mid+1,r,x,y);
	}
	int Que(int p,int l,int r,int ql,int qr) {
		if(!p || (ql<=l && r<=qr)) return s[p];
		int mid=(l+r)>>1,res=0;
		if(ql<=mid) res+=Que(ls[p],l,mid,ql,qr);
		if(qr>mid) res+=Que(rs[p],mid+1,r,ql,qr);
		return res;
	}
} Tree;
int rt[SN][SN],c[N];

void dfs_getans(int u,ll &ans){ // 回答询问，注意要回溯
	c[dep[u]]+=End[u];
	if(u && End[u]) rep(i,1,M) Tree.Upd(rt[i][dep[u]%i],0,Maxlen,dep[u],End[u]);
	for(Node t:Q[u]) if(t.d<=M) ans+=Tree.Que(rt[t.d][t.l%t.d],0,Maxlen,t.l,t.r);
		else for(int i=t.l;i<=t.r;i+=t.d) ans+=c[i];
	for(int v:E[u]) dfs_getans(v,ans);
	c[dep[u]]-=End[u];
	if(u && End[u]) rep(i,1,M) Tree.Upd(rt[i][dep[u]%i],0,Maxlen,dep[u],-End[u]);
}

ll Calc(){
	cmax(Maxlen,10);
	ll ans=0;
	rep(i,1,cnt) ans+=1ll*End[i]*(End[i]-1)/2;
	dfs_getans(0,ans);
	return ans;
}

void Divide(int u){
	vis[u]=1;
	Maxlen=10;
	pre_dfs(u,0,0),M=min(sqrt(Maxlen/10),4.0);
	if(sz[u]==1) return;
	for(auto t:G[u]) if(!vis[t.first]){
		int v=t.first,w=t.second;
		if(!trie[0][w]) trie[0][w]=++cnt;
		dfs(v,u,trie[0][w],1,w,w,w,w,0),Build();
		ll t=Calc();
		ans-=t;
		Clear();
	}
	dfs(u,0,0,0,0,0,0,0,1),Build();
	ll t=Calc();
	Clear();
	ans+=t; Tree.cnt=0;
	rep(i,0,M) rep(j,0,M) rt[i][j]=0;
	for(auto t:G[u]) if(!vis[t.first]) {
		int v=t.first;
		mi=1e9,FindRt(sz[v],v,u);
		Divide(Rt);
	}
}

namespace ptChain{
	int Check(){
		rep(i,2,n-1) if(G[i].size()!=2) return 0;
		return G[1].size()==1 && G[n].size()==1;
	}
	char s[N];
	int cnt;
	ll h1[N],h2[N];
	ll rh1[N],rh2[N];
	void dfs(int u,int f) {
		for(auto t:G[u]) if(t.first!=f) {
			s[++cnt]=t.second;
			dfs(t.first,u);
		}
	}
	int IsPal(int l,int r){
		return 
		(h1[r]-h1[l-1]*Pow1[r-l+1]%P1+P1)%P1==(rh1[l]-rh1[r+1]*Pow1[r-l+1]%P1+P1)%P1 && 
		(h2[r]-h2[l-1]*Pow2[r-l+1]%P2+P2)%P2==(rh2[l]-rh2[r+1]*Pow2[r-l+1]%P2+P2)%P2;
	}

	void Solve() {
		dfs(1,0);
		rep(i,1,cnt) h1[i]=(h1[i-1]*K1+s[i])%P1;
		rep(i,1,cnt) h2[i]=(h2[i-1]*K2+s[i])%P2;
		drep(i,cnt,1) rh1[i]=(rh1[i+1]*K1+s[i])%P1;
		drep(i,cnt,1) rh2[i]=(rh2[i+1]*K2+s[i])%P2;
		ll ans=0;
		rep(i,1,cnt) {
			int l=0,r=min(i-1,cnt-i),res=-1;
			while(l<=r) {
				int mid=(l+r)>>1;
				if(IsPal(i-mid,i+mid)) l=mid+1,res=mid;
				else r=mid-1;
			}
			ans+=res+1;
			if(i>1) {
				l=0,r=min(i-2,cnt-i),res=-1;
				while(l<=r) {
					int mid=(l+r)>>1;
					if(IsPal(i-1-mid,i+mid)) l=mid+1,res=mid;
					else r=mid-1;
				}
				ans+=res+1;
			}
		}
		printf("%lld\n",ans);
	}
}

int main(){
	freopen("string.in","r",stdin),freopen("string.out","w",stdout);
	n=rd();
	rep(i,2,n){
		int u=rd(),v=rd(),w=rd();
		G[u].pb(make_pair(v,w));
		G[v].pb(make_pair(u,w));
	}
	Pow1[0]=Pow2[0]=1;
	rep(i,1,N-1) Pow1[i]=Pow1[i-1]*K1%P1;
	rep(i,1,N-1) Pow2[i]=Pow2[i-1]*K2%P2;
	if(ptChain::Check()) return ptChain::Solve(),0;
	Divide(1);
	printf("%lld\n",ans);
}