[THUWC2018] 城市规划 题解
文章大部分转载 xwd 的博客
这道题目其实出乎意料的简单。
直接让连通块的最上面的节点来进行 dp.
为什么要这样来做呢。
主要是因为这是一棵树。
具有很好的联通关系。
这样的话让每个联通块的最上面的节点记录信息就显得很简洁。
设 f u , j f_{u,j} fu,j 表示考虑点 u u u 的联通块中选择 a u a_u au 和 j j j 这两种颜色,且必选点 u u u,总共的连通块个数。
f u , a u f_{u,a_u} fu,au 表示只选 a u a_u au 这一种颜色。
若当前遍历的边为 ( u , v ) (u,v) (u,v)
若 a u = a v a_u=a_v au=av
f u , i = f u , i × f v , i + f u , a u × f v , i + f u , i × f v , a v + f u , i f_{u,i}=f_{u,i}\times f_{v,i}+f_{u,a_u}\times f_{v,i} +f_{u,i}\times f_{v,a_v}+f_{u,i} fu,i=fu,i×fv,i+fu,au×fv,i+fu,i×fv,av+fu,i
特别的
f u , a u = f u , a u × f v , a v + f u , a u f_{u,a_u}=f_{u,a_u}\times f_{v,a_v}+f_{u,a_u} fu,au=fu,au×fv,av+fu,au
如果 a u ≠ a v a_u\neq a_v au=av
那么 f u , a v = f u , a u × ( f v , a v + f v , a u ) + f u , a v × ( f v , a u + f v , a v + 1 ) f_{u,a_v}=f_{u,a_u}\times (f_{v,a_v}+f_{v,a_u})+f_{u,a_v}\times(f_{v,a_u}+f_{v,a_v}+1) fu,av=fu,au×(fv,av+fv,au)+fu,av×(fv,au+fv,av+1)
其他项没有必要更新。因为选了这个点的话本来就有两种不同的颜色了。
稍微解释一下为什么有 f u , a u × f v , a u f_{u,a_u}\times f_{v,a_u} fu,au×fv,au 这一项。
因为 a v a_v av 已经在 v v v 这里必选了。相当于 f v , a u f_{v,a_u} fv,au 本身就带有 a u a_u au 和 a v a_v av 这两种颜色。
感觉有点废话。。
这样的话暴力是 O ( n 2 ) O(n^2) O(n2) 的。非常不美妙。
但是这个比较显然就直接动态开点权值线段树解决了。
然后对于 a u = a v a_u=a_v au=av 的时候需要用线段树合并。
对于 a u ≠ a v a_u\neq a_v au=av 的部分就是用一个单点查询和修改即可。
答案就是根节点的线段树维护的和。
代码还没有写。。。
#include<bits/stdc++.h>
#define int long long
using namespace std;
inline int rd(){
int s=0,w=1;char ch=getchar();
for(;ch<'0'||ch>'9';ch=getchar())if(ch=='-')w=-1;
for(;ch>='0'&&ch<='9';ch=getchar())s=s*10+ch-'0';
return s*w;
}
inline void wr(int x){
if(x<0)x=-x,putchar('-');
if(x>9)wr(x/10);
putchar(x%10+'0');
}
const int N=1e5+5,M=80,p=998244353;
int n;
int tot,ans;
int a[N],id[N];
struct node{
int l,r;
int sum,laz;
}tr[N*M];
vector<int> G[N];
inline void add(int u,int v){G[u].push_back(v);}
inline void New(int &rt){
rt=++tot;
tr[rt].laz=1;
}
inline void pushup(int k){
tr[k].sum=(tr[tr[k].l].sum+tr[tr[k].r].sum)%p;
}
inline void pushdown(int k){
if(tr[k].laz==1)return;
int ls=tr[k].l,rs=tr[k].r,vl=tr[k].laz;
if(ls){
tr[ls].laz=tr[ls].laz*vl%p;
tr[ls].sum=tr[ls].sum*vl%p;
}
if(rs){
tr[rs].laz=tr[rs].laz*vl%p;
tr[rs].sum=tr[rs].sum*vl%p;
}
tr[k].laz=1;
}
inline void ins(int &rt,int l,int r,int x){//新建
if(!rt)New(rt);
if(l==r){
tr[rt].sum=1;return;
}
int mid=(l+r)>>1;
if(x<=mid)ins(tr[rt].l,l,mid,x);
else ins(tr[rt].r,mid+1,r,x);
pushup(rt);
}
inline void upd(int &rt,int l,int r,int x,int fu_u,int fv_u,int fv_v){//
if(!rt)New(rt);
if(l==r){
tr[rt].sum=(tr[rt].sum*(fv_u+fv_v+1)%p+fu_u*(fv_v+fv_u))%p;
return;
}
int mid=(l+r)>>1;
pushdown(rt);
if(x<=mid)upd(tr[rt].l,l,mid,x,fu_u,fv_u,fv_v);
else upd(tr[rt].r,mid+1,r,x,fu_u,fv_u,fv_v);
pushup(rt);
}
inline int query(int rt,int l,int r,int x){
if(!rt)return 0;
if(l==r)return tr[rt].sum;
pushdown(rt);
int mid=(l+r)>>1;
if(x<=mid)return query(tr[rt].l,l,mid,x);
else return query(tr[rt].r,mid+1,r,x);
}
inline int merge(int x,int y,int l,int r,int fu_u,int fv_v,int A){
if(!x&&!y)return 0;
if(!x){
tr[y].sum=(tr[y].sum*fu_u)%p;
tr[y].laz=(tr[y].laz*fu_u)%p;
return y;
}
if(!y){
tr[x].sum=(tr[x].sum*(fv_v+1))%p;
tr[x].laz=(tr[x].laz*(fv_v+1))%p;
return x;
}
if(l==r){//f[u][a[u]] 肯定在上面或者下面。
if(l==A)tr[x].sum=tr[x].sum*(fv_v+1)%p;
else tr[x].sum=(tr[x].sum*tr[y].sum%p+fu_u*tr[y].sum%p+tr[x].sum*fv_v%p+tr[x].sum)%p;
return x;
}
int mid=(l+r)>>1;
pushdown(x);pushdown(y);
int c=0;New(c);
tr[c].l=merge(tr[x].l,tr[y].l,l,mid,fu_u,fv_v,A);
tr[c].r=merge(tr[x].r,tr[y].r,mid+1,r,fu_u,fv_v,A);
pushup(c);
return c;
}
inline void dfs(int u,int FA){
ins(id[u],1,n,a[u]);
for(int v:G[u]){
if(v==FA)continue;
dfs(v,u);
int fu_u=query(id[u],1,n,a[u]);
int fv_v=query(id[v],1,n,a[v]);
if(a[u]==a[v]){
id[u]=merge(id[u],id[v],1,n,fu_u,fv_v,a[u]);
}
else{
int fv_u=query(id[v],1,n,a[u]);
upd(id[u],1,n,a[v],fu_u,fv_u,fv_v);
}
}
ans+=tr[id[u]].sum;ans%=p;
}
signed main(){
n=rd();
for(int i=1;i<=n;++i)a[i]=rd();
for(int i=1;i<n;++i){
int x=rd(),y=rd();
add(x,y);add(y,x);
}
dfs(1,0);
wr(ans);
return 0;
}