题解：P11149 [THUWC 2018] 城市规划

一种复杂度很大但是常数小到可以过题的做法。

题意：给你一颗有颜色的树，求有多少个联通块满足内部至多只有两种颜色。

到这个题要我们求的是符合要求的联通块个数，自然而然的想到了点分治，那么问题转化为了必须包含根节点的联通块个数。

于是联通块必须包含根节点所在的颜色。

我们可以先把根节点所在的且所有点的颜色都和根节点相同的最大联通块 S 找出来。

这一个联通块里的点是无论另一种颜色怎么选择都一定可以被选到的。

而与这个联通块相邻的那些点是比较特殊的，因为如果要选择这个点子树内的点，就一定会选择这个点，所以对于这些点的子树而言，联通块所包含的两种颜色是已经确定了的。

这意味着，我们可以把所有非 S 内的点的信息全部都挂到一个和相邻的点上。

然后我们就可以考虑如何 DP 了。

假设已经确定了哪两种颜色是可行的，那么是简单的，假设定义 f(x) 代表以 x 为根的子树的方案数。

转移就是 f(x)=\prod\limits_{y\in son(x)}\bigg(f(y)+1\bigg)

对于不同的选颜色方案，可以选择的点是不同的。

不过根据前文可以得到如下结论。

1. 无论如何，S 内的点都可以被选到。

2. 不在 S 内的点，信息都可以挂到一个与 S 相邻的点上。

于是做法也顺理成章的出来了：

我们先把所有和 S 相邻的点的 f 值求出来，

再把 S 这颗树的 f 值求出来，

接着按颜色枚举与 S 相邻的点的颜色，

那么对于每一种颜色，其对于 S 这棵树的改变相当于给一些节点下面挂一个 f 值，于是我们使用 DDP 来维护这个东西。

于是这个题就做完了，不过要注意的是，对于每一种颜色，我们求出的方案数是根节点颜色加当前颜色的所有方案的数量，因此还需要减去 S 内部产生方案的数量。

都做这种题了应该都会 DDP 吧。

复杂度是 O(n\log^3n)，不过常数极小，在一个不算惊险的时间跑过去了，一些实现上和DDP的细节可以看代码。

code:

#include <bits/stdc++.h>
#define ll long long
#define pii pair<int,int>
using namespace std;
const int N=5e5+10,M=998244353;
int n,ans,res,old_res,rt;
int a[N];
void do_mod(int &x){
    if(x>=M)x-=M;
}
vector<int>e[N];
int ksm(int x,int y){
    int res=1;
    while(y)
    {
        if(y&1)res=1ll*res*x%M;
        x=1ll*x*x%M;
        y>>=1;
    }
    return res;
}
struct barycenter/*求重心单独开结构体防止变量名重名*/{
    int siz[N],msiz[N];
    int minn=n,id,sum;
    void dfs(int x,int fd){
        siz[x]=1;
        msiz[x]=0;
        for(auto y:e[x])
        {
            if(y==fd)continue;
            dfs(y,x);
            siz[x]+=siz[y];
            msiz[x]=max(msiz[x],siz[y]);
        }
//      msiz[x]=max(siz[x],sum-siz[x]-1);
        msiz[x]=max(msiz[x],sum-siz[x]);
        if(msiz[x]<minn)
        {
            minn=msiz[x];
            id=x;
        }
    }
    void rebuild(int &x,int sig){
        sum=sig;
        minn=sig;
        id=0;
        dfs(x,sig);
        x=id;
    }
    void ddfs(int x,int fd){
        siz[x]=1;
        for(auto y:e[x])
        {
            if(y==fd)continue;
            dfs(y,x);
            siz[x]+=siz[y];
        }
    }
}B;
int deep[N],siz[N],fa[N],son[N];//描述树性质的数组 
int dfn[N],id[N],top[N],cnt,len[N];//描述树链性质的数组
int g[N],f[N];//dp数组 
int old_g[N];//为了撤销用 
pii point[N];int poi;//所有的与S相邻的点  
int croot,ccol;//对于某个与S相邻的点，其子树内可以选择的两种颜色 
void dfs_dp(int x,int fd)/*预处理的dfs*/{
    f[x]=1;
    g[x]=1;
    siz[x]=1;
    son[x]=0;
    fa[x]=fd;
    deep[x]=deep[fa[x]]+1;
    top[x]=0;
    for(int y:e[x])
    {
        if(y==fd)continue;
        if(a[y]!=a[x])
        {
            point[++poi]={y,x};
            continue;
        }
        dfs_dp(y,x);
        siz[x]+=siz[y];
        if(siz[y]>siz[son[x]])
        {
            son[x]=y;
        }
        f[x]=1ll*f[x]*(f[y]+1)%M;
    }
    for(int y:e[x])
    {
        if(y==fd)continue;
        if(y==son[x]||a[y]!=a[x])continue;
        g[x]=1ll*g[x]*(f[y]+1)%M;
    }
}
void dfs_dfn(int x){
    dfn[x]=++cnt;
    id[cnt]=x;
    len[x]=0;
    if(!top[x])top[x]=x;
    len[top[x]]++;
    if(son[x])
    {
        top[son[x]]=top[x];
        dfs_dfn(son[x]);
    }
    for(int y:e[x])
    {
        if(y==fa[x]||y==son[x]||a[y]!=a[x])continue;
        dfs_dfn(y);
    }
}
void dfs_y(int x,int fd){
    f[x]=1;
    for(int y:e[x])
    {
        if(y==fd)continue;
        if(a[y]!=croot&&a[y]!=ccol)continue;
        dfs_y(y,x);
        f[x]=1ll*f[x]*(f[y]+1)%M;
    }
}
struct mac/*矩阵*/{
    int a[2][2];
}empty,ifg[N*4],*root[N],*tree;
//为了减小常数，每一个链开一个线段树 
mac *inc;
int add;
mac operator *(mac x,mac y){
    mac res=empty;
    for(int i=0;i<2;i++)
    {
        for(int k=0;k<2;k++)
        {
            for(int j=0;j<2;j++)
            {
                res.a[i][j]=(res.a[i][j]+1ll*x.a[i][k]*y.a[k][j])%M;
            }
        }
    }
    return res;
}
struct seg{
    void build(int l,int r,int p){
        if(l==r)
        {
            tree[p].a[0][0]=g[id[add+l]];
            tree[p].a[1][0]=g[id[add+r]];
            tree[p].a[1][1]=1;
            return;
        }
        int mid=(l+r)/2;
        build(l,mid,p*2);
        build(mid+1,r,p*2+1);
        tree[p]=tree[p*2+1]*tree[p*2];
    }
    void change(int pos,int num,int l,int r,int p){
        if(l==r)
        {
            tree[p].a[0][0]=tree[p].a[1][0]=num;
            return;
        }
        int mid=(l+r)/2;
        if(mid>=pos)
        {
            change(pos,num,l,mid,p*2);
        }
        else
        {
            change(pos,num,mid+1,r,p*2+1);
        }
        tree[p]=tree[p*2+1]*tree[p*2];
    }
}T;
void dfs_build(int x,int fd){
    if(top[x]==x)
    {
        root[x]=inc;
        inc+=len[x]*4;
        add=dfn[x]-1;
        tree=root[x];
        T.build(1,len[x],1);
    }
    for(int y:e[x])
    {
        if(y==fd||a[y]!=a[x])continue;
        dfs_build(y,x);
    }
}
bool cmp(pii x,pii y){
    return a[x.first]<a[y.first];
}
void change(int x,int num)/*将g[x]的值*num */{
    if(old_g[x]==-1)old_g[x]=g[x];
    g[x]=1ll*g[x]*num%M;
    while(x)
    {
        add=dfn[top[x]]-1;
        tree=root[top[x]];
        int ni=ksm(tree[1].a[1][0]+1,M-2);
        T.change(dfn[x]-dfn[top[x]]+1,g[x],1,len[top[x]],1);
        int s=tree[1].a[1][0];
        if(top[x]==rt)
        {
            res=s;
            break;
        }
        x=fa[top[x]];
        int new_f=s;
        if(old_g[x]==-1)
        {
            old_g[x]=g[x];
        }
        g[x]=1ll*g[x]*ni%M*(new_f+1)%M;
    }
}
void back(int x)/*DDP的撤销*/{
    if(~old_g[x])
    {
        g[x]=old_g[x],old_g[x]=-1;
    }
    while(x)
    {
        add=dfn[top[x]]-1;
        tree=root[top[x]];
        T.change(dfn[x]-dfn[top[x]]+1,g[x],1,len[top[x]],1);
        if(top[x]==rt)
        {
            break;
        }
        x=fa[top[x]];
        if(~old_g[x])
        {
            g[x]=old_g[x];
            old_g[x]=-1;
        }
    }
}
void solve(int x,int sum)/*点分治的主要函数*/{ 
    if(sum==1)
    {
        ans++;
        return;
    }
    if(sum==2)
    {
        ans+=3;
        return;
    }
    B.rebuild(x,sum); 
    B.ddfs(x,sum);//找到重心
    poi=0,cnt=0,inc=ifg,rt=x;
    dfs_dp(x,0);//求出S这颗子树的f值,顺便预处理DDP的树剖 
    do_mod(ans+=f[x]),old_res=f[x];
    dfs_dfn(x);//DDP的树剖预处理 
    for(int i=1;i<=poi;i++)
    {
        croot=a[x],ccol=a[point[i].first];
        dfs_y(point[i].first,point[i].second);//找到与S相邻的点的f值 
    }
    dfs_build(x,0);//DDP的建线段树 
    sort(point+1,point+1+poi,cmp);
    for(int l=1;l<=poi;l++)
    {
        int r=l;
        while(r<poi&&a[point[r+1].first]==a[point[l].first])r++;
        //颜色相同的块一起处理 
        for(int i=l;i<=r;i++)
        {
            change(point[i].second,f[point[i].first]+1);
        }
        do_mod(ans+=(res-old_res)<0?(res-old_res+M):res-old_res);
        //要减去S内部的贡献 
        for(int i=l;i<=r;i++)
        {
            back(point[i].second);
        }l=r;
    }
//  cout<<x<<' '<<ans<<'\n';
    for(int y:e[x])
    {
        for(int o=0;o<(int)e[y].size()-1;o++)
        {
            if(e[y][o]==x)
            {
                swap(e[y][o],e[y][(int)e[y].size()-1]);
            }
        }
        e[y].pop_back();
        solve(y,B.siz[y]);
    }
}
signed main()
{
//  freopen("in.in","r",stdin);
//  freopen("1.out","w",stdout);
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    //
    cin>>n;
    for(int i=1;i<=n;i++)
    {
        cin>>a[i];
    }
    for(int i=1;i<n;i++)
    {
        int u,v;
        cin>>u>>v;
        e[u].push_back(v);
        e[v].push_back(u);
    }
    memset(old_g,-1,sizeof old_g);
    solve(1,n);
    cout<<ans;
    return 0;
}
/*
考虑DDP部分的细节
f[i]代表答案,g[i]代表轻儿子的答案
转移为
f[i]=g[i]*(f[son]+1)
换成矩阵的形式 
[f[son],1]*U=[f[i],1]
于是转移矩阵 U就是 
g[i],0
g[i],1

而DDP的撤销部分是用一种记录的方式实现的
在树剖对g值进行改变时,记录下原本的g值
然后做完之后再复原 

*/