CF1916E Happy Life in University

tags

枚举 线段树

中文题面

给定 $n$ 个点、点 $u$ 有颜色 $a_u$ 的树，定义 $f(u,v)$ 为 $u\to v$ 的路径上的颜色数目，$u$ 不是 $v$ 祖先时无定义。对于所有点对 $(u,v)$，其中 $u$ 不一定相异于 $v$，设 $l$ 为 $u,v$ 的最近公共祖先，求 max ⁡ { f ( l , u ) ⋅ f ( l , v ) } \max\{ f(l, u) \cdot f(l, v) \} max{f(l,u)⋅f(l,v)}。

多测，$1\le \sum n\le 3\cdot 10^5,1\le a_i\le n,1\le p_i\le i-1$。

第一行输入 $t$ 表示数据组数。对于每组数据，第一行输入 $n$，第二行输入 $n-1$ 个数 $p_2\sim p_n$，$p_i$ 为 $i$ 的父亲的编号，第三行输入 $n$ 个数 $a_1\sim a_n$ 表示编号为 $i$ 的点的颜色。

对于每组数据，输出一行一个数字，表示 max ⁡ { f ( l , u ) ⋅ f ( l , v ) } \max\{ f(l, u) \cdot f(l, v) \} max{f(l,u)⋅f(l,v)}。

思路

考虑枚举 $lca$，找出每个子树中叶子到 $lca$ 路径颜色数的最大值 $[max]$、次大值 $[min]$，则该子树答案为

$ans=[max]\ast [min]$

问题转化为动态维护路径的颜色数最大值，可以用线段树来解决。

从叶子节点出发，每次直接给子树 $+1$，注意在算重的位置减去贡献即可。

代码

#include<bits/stdc++.h>
using namespace std;
#define int long long
#define A(x,y) x=(x+y)%mod
#define N 3e5 + 5
const int M = N<<2, mod = 998244353;
int T,n,c[N],p[N],dfn[N],low[N],tot, asw;vector<int>v[N],s[N];
struct tree{
	int l[M],r[M],tag[M],ma[M];
	#define ls (k<<1)
	#define rs (k<<1|1)
	void Pd(int k,int v){
		tag[k]+=v;
		ma[k]+=v;
	}
	void pd(int k){
		tag[k]&&(Pd(ls,tag[k]),Pd(rs,tag[k]),tag[k]=0);
	}
	void pu(int k){
		ma[k]=max(ma[ls],ma[rs]);
	}
	void build(int L=1,int R=n,int k=1){
		l[k]=L,r[k]=R,tag[k]=ma[k]=0;
		if(L==R)return;
		int mid=L+R>>1;
		build(L,mid,ls),build(mid+1,R,rs);
	}
	void add(int L,int R,int v,int k=1){
		if(R<l[k]||r[k]<L)return;
		if(L<=l[k]&&r[k]<=R)return Pd(k,v);
		pd(k),add(L,R,v,ls),add(L,R,v,rs),pu(k);
	}
	int ask(int L,int R,int k=1){
		if(R<l[k]||r[k]<L)return 0;
		if(L<=l[k]&&r[k]<=R)return ma[k];
		return pd(k),max(ask(L,R,ls),ask(L,R,rs));
	}
}t;
void dfs(int u){int tmp=p[c[u]];s[tmp].push_back(u),p[c[u]]=u,dfn[u]=++tot;
	for(auto i:v[u])dfs(i);p[c[u]]=tmp,low[u]=tot;
}
void DFS(int u){
	for(auto i:v[u])DFS(i);int f1=1,f2=1;t.add(dfn[u],low[u],1);
	for(auto i:s[u])t.add(dfn[i],low[i],-1);
	for(auto i:v[u]){int f=t.ask(dfn[i],low[i]);
		if(f>f1)f2=f1,f1=f;else f2=max(f2,f);
	}
	asw=max(asw,1ll*f1*f2);
}
signed main(){ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
	cin>>T;
	while(T--){
		cin>>n,asw=tot=0;
		for(int x,i=2;i<=n;i++)cin>>x,v[x].push_back(i);
		for(int i=1;i<=n;i++)cin>>c[i];
		dfs(1),t.build(),DFS(1);
		cout<<asw<<"\n";
		for(int i=0;i<=n;i++)v[i].clear(),s[i].clear();
	}
}