P7422 「PMOI-2」城市 题解

题目描述

给定一张 \(n\) 个点， \(m\) 条边的无向连通图，点有颜色 \(c_i\) 。

若 \(A\to 1\) 的任意简单路径必须经过 \(B\) ，则称 \(B\) 在 \(A\to 1\) 的必经之路上。

若 \(B\to A\) 的任意简单路径与 \(C\to A\) 的任意简单路径没有公共边，则称 \(B,C\) 关于 \(A\) 互不影响。

记 \(f(A,k)\) 为满足下列所有条件的 \(k\) 元集合 \(\{B_1,\cdots,B_k\}\) 个数：

• \(\forall 1\le i\le k\) ， \(B_i\) 与 \(A\) 的颜色不同且 \(A\) 在 \(B_i\to 1\) 的必经之路上。
• \(\forall 1\le i\lt j\le k\) ， \(B_i,B_j\) 颜色相同且 \(B_i,B_j\) 关于 \(A\) 互不影响。

给定 \(K\) ，求 \(\sum_{i=1}^n\sum_{k=1}^Kf(i,k)\) ，对 \(998244353\) 取模。

数据范围

• \(1\le n\le 5\cdot 10^5,1\le m\le\min(10^6,\frac{n(n-1)}2),1\le K\le 20,1\le c_i\le 10^9\) 。

时间限制 \(\texttt{2.5s}\) ，空间限制 \(\texttt{500MB}\) 。

分析

建出圆方树，钦定 \(1\) 号点为根，条件变为：

• \(\{B_1,.\cdots,B_k\}\) 颜色相同且与 \(A\) 不同。
• \(B_i\) 在 \(A\) 的不同子树中。

枚举 \(B_i\) 颜色，记 \(cnt_u\) 为 \(u\) 子树中该颜色出现次数， \(A\) 的贡献为：

\[\sum_{k=1}^K[x^k]\prod_{u\in son(A)}(1+cnt_ux)\\ \]

当 \(k\ge 2\) 时， \(A\) 一定出现在由该颜色的点构成的虚树上。对每种颜色建虚树，对虚树上每个点维护 \(\prod\limits_{u\in son(A)}(1+cnt_ux)\) 的低 \(K+1\) 位，注意只有圆点的贡献才需要计入答案。

当 \(k=1\) 时， \(A\) 的贡献为子树内与 \(A\) 不同色的点个数。 \(dfs\) 时维护每种颜色出现次数，用进入 \(A\) 和回溯前的数据作差即可。

时间复杂度 \(\mathcal O(n(\log n+K))\) 。

#include<bits/stdc++.h>
using namespace std;
const int maxn=1e6+5,mod=998244353;
int k,m,n,u,v,cnt,num;
int col[maxn],dfn[maxn],low[maxn];
stack<int> st;
vector<int> g[maxn],vec[maxn];
namespace tree
{
    int x,cnt,res,top,tot;
    int c[maxn],d[maxn],dfn[maxn],sz[maxn];
    int st[maxn];
    int fa[maxn][18];
    int head[maxn],to[maxn],nxt[maxn];
    vector<int> g[maxn];
    void dfs1(int u)
    {
        dfn[u]=++cnt,sz[u]=u<=n;
        int tmp=c[col[u]]++;
        for(auto v:g[u])
        {
            d[v]=d[u]+1,fa[v][0]=u;
            for(int i=1;i<18;i++) fa[v][i]=fa[fa[v][i-1]][i-1];
            dfs1(v),sz[u]+=sz[v];
        }
        if(u<=n) res=(res+sz[u]-(c[col[u]]-tmp))%mod;
    }
    int lca(int u,int v)
    {
        if(d[u]<d[v]) swap(u,v);
        for(int i=17;i>=0;i--) if(d[fa[u][i]]>=d[v]) u=fa[u][i];
        if(u==v) return u;
        for(int i=17;i>=0;i--) if(fa[u][i]!=fa[v][i]) u=fa[u][i],v=fa[v][i];
        return fa[u][0];
    }
    void addedge(int u,int v)
    {
        nxt[++tot]=head[u],to[tot]=v,head[u]=tot;
    }
    void dfs2(int u)
    {
        c[u]=col[u]==x;
        vector<int> f(k+1); f[0]=1;
        for(int i=head[u];i;i=nxt[i])
        {
            int v=to[i];
            dfs2(v),c[u]+=c[v];
            for(int j=k;j>=1;j--) f[j]=(f[j]+1ll*f[j-1]*c[v])%mod;
        }
        if(u<=n&&col[u]!=x) for(int j=2;j<=k;j++) res=(res+f[j])%mod;
    }
    void solve()
    {
        memcpy(c+1,col+1,4*n);
        sort(c+1,c+n+1);
        int cnt=unique(c+1,c+n+1)-c-1;
        for(int i=1;i<=n;i++) vec[col[i]=lower_bound(c+1,c+cnt+1,col[i])-c].push_back(i);
        memset(c,0,sizeof(c));
        d[1]=1,dfs1(1);
        for(x=1;x<=cnt;x++)
        {
            sort(vec[x].begin(),vec[x].end(),[&](int a,int b){return dfn[a]<dfn[b];});
            head[1]=0,st[top=1]=1;
            for(auto u:vec[x])
            {
                if(u==1) continue;
                int p=lca(u,st[top]);
                if(p!=st[top])
                {
                    while(dfn[p]<dfn[st[top-1]]) addedge(st[top-1],st[top]),top--;
                    if(dfn[p]>dfn[st[top-1]]) head[p]=0,addedge(p,st[top--]),st[++top]=p;
                    else addedge(st[top-1],st[top]),top--;
                }
                head[u]=0,st[++top]=u;
            }
            while(top>=2) addedge(st[top-1],st[top]),top--;
            dfs2(1);
        }
        printf("%d\n",res);
    }
}
void tarjan(int u)
{
    dfn[u]=low[u]=++cnt,st.push(u);
    for(auto v:g[u])
    {
        if(!dfn[v])
        {
            tarjan(v);
            low[u]=min(low[u],low[v]);
            if(low[v]>=dfn[u])
            {
                tree::g[u].push_back(++num);
                static int p=0;
                do p=st.top(),st.pop(),tree::g[num].push_back(p);
                while(p!=v);
            }
        }
        else low[u]=min(low[u],dfn[v]);
    }
}
int main()
{
    scanf("%d%d%d",&n,&m,&k),num=n;
    for(int i=1;i<=n;i++) scanf("%d",&col[i]);
    for(int i=1;i<=m;i++) scanf("%d%d",&u,&v),g[u].push_back(v),g[v].push_back(u);
    tarjan(1);
    tree::solve();
    return 0;
}