[THUWC2018]城市规划的题解

[THUWC2018]城市规划

连通块问题，我们考虑树形 DP。

设 $f_{u,j}$ 表示在以 $u$ 为根的子树内，选的颜色集合为 $a_u,j$（两个颜色都必须选）且必须选点 $u$ 的情况下的连通块个数。

特殊的，$f_{u,a_u}$ 表示颜色只有 $a_u$ 的情况数。

对于 $v\in so{n}_u$，有：

---

若 $a_u=a_v$：

对于每一个颜色 $i\in [1,n]$，有：

$$f_{u,i}=f_{v,i}\times f_{u,a_u}+f_{v,a_v}\times f_{u,i}+f_{u,i}\times f_{v,i}+f_{u,i}$$

特别的，有：

$$f_{u,a_u}=f_{u,a_u}\times (f_{v,a_v}+1)$$

---

否则，当$a_u\ne a_v$时：

$$f_{u,a_v}=f_{u,a_v}\times (f_{v,a_v}+f_{v,a_u}+1)+f_{u,a_u}\times (f_{v,a_v}+f_{v,a_u})$$

---

但是，这样子时间复杂度太大，接受不了。

我们发现上面对于的操作分别为线段树合并和线段树单点查询、修改，以此优化即可。

注意代码细节。

代码：

#include<bits/stdc++.h>
#define ll long long
using namespace std;
 
const int N=1e5+5;
const ll mod=998244353;
int n,a[N];
struct node{
	int l,r;
	ll sum,mul;
}tree[N<<6];
int root[N],tot;
void clean(int k){
	tree[k].sum%=mod,tree[k].mul%=mod;
}
int build(){
	int k=++tot;
	tree[k].mul=1;
	return k;
}
void push_down(int k){
	if(tree[k].mul==1) return;
	if(tree[k].l){
		tree[tree[k].l].mul*=tree[k].mul;
		tree[tree[k].l].sum*=tree[k].mul;
		clean(tree[k].l);
	}
	if(tree[k].r){
		tree[tree[k].r].mul*=tree[k].mul;
		tree[tree[k].r].sum*=tree[k].mul;
		clean(tree[k].r);
	}
	tree[k].mul=1;
}
void push_up(int k){
	tree[k].sum=tree[tree[k].l].sum+tree[tree[k].r].sum;
	clean(k);
}
ll query(int k,int l,int r,int pos){
	if(!k) return 0;
	if(l==r) return tree[k].sum%mod;
	push_down(k);
	int mid=(l+r)>>1;
	if(pos<=mid) return query(tree[k].l,l,mid,pos);
	else return query(tree[k].r,mid+1,r,pos);
}
 
int insert(int k,int l,int r,int pos){
	if(!k) k=build();
	if(l==r){
		tree[k].sum=1;
		return k;
	}
	push_down(k);
	int mid=(l+r)>>1;
	if(pos<=mid) tree[k].l=insert(tree[k].l,l,mid,pos);
	else tree[k].r=insert(tree[k].r,mid+1,r,pos);
	push_up(k);
	return k;
}
int merge(int x,int y,int l,int r,ll tu,ll tv,int cu){
	if(!x&&!y) return 0;
	if(!x){//f[u][l-r]都没有 
		tree[y].sum*=tu;
		tree[y].mul*=tu;
		clean(y);
		return y;
	}
	if(!y){
		tree[x].sum*=(tv+1ll);
		tree[x].mul*=(tv+1ll);
		clean(x);
		return x;
	}
	if(l==r){
		if(l==cu) tree[x].sum=tree[x].sum*(tv+1)%mod;
		else tree[x].sum=tree[y].sum*tu%mod+tree[x].sum*tv%mod+tree[x].sum*tree[y].sum%mod+tree[x].sum;
		tree[x].sum%=mod;
		return x;
	}
	int mid=(l+r)>>1;
	push_down(x),push_down(y);
	tree[x].l=merge(tree[x].l,tree[y].l,l,mid,tu,tv,cu);
	tree[x].r=merge(tree[x].r,tree[y].r,mid+1,r,tu,tv,cu);
	push_up(x);
	return x;
}
int mul(int k,int l,int r,int pos,ll tu,ll tv,ll fv_cu){
	if(!k) k=build();
	if(l==r){
		tree[k].sum=tree[k].sum*(fv_cu+tv+1)%mod+tu*(fv_cu+tv)%mod;
		tree[k].sum%=mod;
		return k;
	}
	push_down(k);
	int mid=(l+r)>>1;
	if(pos<=mid) tree[k].l=mul(tree[k].l,l,mid,pos,tu,tv,fv_cu);
	else tree[k].r=mul(tree[k].r,mid+1,r,pos,tu,tv,fv_cu);
	push_up(k);
	return k;
}
ll ans=0;
vector<int> G[N];
void dfs(int u,int fa){
	root[u]=insert(root[u],1,n,a[u]);
	ll tu,tv,fv_cu;
	for(auto v : G[u]){
		if(v==fa) continue;
		dfs(v,u);
		tu=query(root[u],1,n,a[u]);tv=query(root[v],1,n,a[v]);
		if(a[v]!=a[u]){
			fv_cu=query(root[v],1,n,a[u]);
			root[u]=mul(root[u],1,n,a[v],tu,tv,fv_cu);
		}
		else{
			root[u]=merge(root[u],root[v],1,n,tu,tv,a[u]);
		}
	}
	ll sum=tree[root[u]].sum;
	ans=(ans+sum)%mod;
}
 
int main(){
	scanf("%d",&n);
	for(int i=1;i<=n;i++)scanf("%d",&a[i]);
	int x,y;
	for(int i=1;i<n;i++){
		scanf("%d%d",&x,&y);
		G[x].push_back(y);
		G[y].push_back(x);
	}
	dfs(1,0);
	printf("%lld\n",ans);
	return 0;
}