题解：P11194 [COTS 2021] 县 Županije

· · 题解

首先我们发现每个县一定是一个连通块，否则显然无解。

对于连接不同县的两个点的边 (u,v)，设其所属的县城分别为 rt_u,rt_v，那么根据题意有：

\begin{cases}dis(u,rt_u)<dis(u,rt_v)\\dis(v,rt_v)<dis(v,rt_u)\\dis(u,rt_v)=dis(v,rt_v)+1\\dis(v,rt_u)=dis(u,rt_u)+1\end{cases}

其中 dis(x,y) 表示 x,y 间简单路径长度。于是不难解出 dis(u,rt_u)=dis(v,rt_v)，发现这一条件是充要的，于是我们只要选出若干县城使其符合这一结构即可。

显然缩点后县之间也构成一棵树 T'。我们令 f_u=0/1 表示只考虑点 u 所在县在 T' 上的子树，u 能否成为县城，尝试进行自底向上的转移。对于某一县 i 内的一个点 u 且 f_u=1，假设该县在 T' 上的父亲为 j，其中关键边为 (x,y)（不妨 x 在县 i 内），那么 u 的贡献是县 j 中所有与 y 的距离为 dis(u,x) 的点。

如上图，若红色县内 f_9=1 也就是说 9 可能成为县城，那么其对于黑色县的贡献就是点 2,6 有可能成为县城，因为它们都满足 dis(2,3)=dis(6,3)=dis(9,5)=2。

而这种与“树上与某一个点距离相同的一些点”相关的问题，比较自然地想到点分治维护。所以我们对每一个县进行点分治，按照上面的方式向上转移，一个点能够成为县城当且仅当其每个子连通块均存在县城能够让它有可能成为县城（可能有点绕），于是计算出所有 f_u 的值。考虑如何输出一组可行的解，可以自上而下遍历县并枚举县中的所有点判断是否合法即可。

思路不算繁琐，但是代码实现比较复杂（也有可能是我太菜了 qaq），可以结合代码理解。使用 \mathcal{O}(1) 的 LCA 可以做到 \mathcal{O(n\log n)} 的复杂度，关于 \mathcal{O}(1) 的 LCA 可以参考这篇博客。

::::info[一点有可能用得上的贴士？]

::::

%:include <bits/stdc++.h>
using namespace std;

#define il inline
// 快读

#define eb emplace_back
const int N=2e5+5;
int n,m,bel[N];
vector<int> edg[N],se[N];

const int Lg=20;
int top[N],dep[N],dfn[N],g[Lg][N],ft[N],tim;
void pre(int u,int fa) {
    top[u]=(bel[u]^bel[fa])?u:top[fa],g[0][dfn[u]=++tim]=ft[u]=fa;
    for(int v:edg[u]) if(v^fa) dep[v]=dep[u]+1,pre(v,u);
}
#define get(x,y) (dfn[x]<dfn[y]?x:y)
il void init() {
    for(int i=1;(1<<i)<=n;i++)
        for(int j=1;j+(1<<i)-1<=n;j++)
            g[i][j]=get(g[i-1][j],g[i-1][j+(1<<i-1)]);
}
il int lca(int u,int v) {
    if(u==v) return u;
    if((u=dfn[u])>(v=dfn[v])) swap(u,v);
    int k=__lg(v-u++);
    return get(g[k][u],g[k][v-(1<<k)+1]);
}
il int dis(int u,int v) {return dep[u]+dep[v]-(dep[lca(u,v)]<<1);}

#define dmx(x,y) (x<y)&&(x=y)
int vis[N],tp[N],sz[N],rdep[N],mxd[N];
int fnd(int u,int fa,int siz) {
    for(int v:edg[u]) {
        if(vis[v]||v==fa||bel[u]^bel[v]) continue ;
        if(sz[v]>=(siz>>1)) return fnd(v,u,siz);
    }
    return u;
}
void dfs(int u,int rt,int fa) {
    mxd[rt]=max(mxd[rt],rdep[u]),sz[u]=1;
    for(int v:edg[u]) {
        if(vis[v]||v==fa||bel[u]^bel[v]) continue ;
        rdep[v]=rdep[u]+1,dfs(v,rt,u),sz[u]+=sz[v];
    }
}
void dfz(int u) {
    vis[u]=1,rdep[u]=0,dfs(u,u,0);
    for(int v:edg[u]) {
        if(vis[v]||bel[u]^bel[v]) continue ;
        int vrt=fnd(v,0,sz[v]); dfz(vrt),tp[vrt]=u;
    }
}

int f[N],vid[N];
vector<int> son[N],cnt[N],tnc[N],tg;
il void upd(int u,int d) {
    if(d<=mxd[u]) ++cnt[u][d];
    for(int v=u,o=u;u=tp[u];v=u) {
        int k=d-dis(o,u);
        if(k>=0&&k<=mxd[u]) ++cnt[u][k],++tnc[v][k];
    }
}
il int ask(int u) {
    int tmp=cnt[u][0],o=u;
    for(int v=u;u=tp[u];v=u) {
        int k=dis(o,u);
        if(k<=mxd[u]) tmp+=cnt[u][k]-tnc[v][k];
    }
    return tmp;
}
void solve(int u) {
    int loc=ft[top[se[u][0]]];
    for(int v:son[u]) solve(v);
    for(int p:se[u]) {
        f[p]=(ask(p)==son[u].size());
        if(f[p]&&loc) {
            int d=dis(p,loc)-1;
            if(!vid[d]) upd(loc,d),vid[d]=1,tg.eb(d);
        }
    }
    for(int d:tg) vid[d]=0; tg.clear();
}

int ans[N];
void getans(int u) {
    ans[bel[u]]=u;
    for(int p:son[bel[u]]) {
        int k=top[se[p][0]];
        for(int v:se[p]) if(f[v]&&dis(u,k)==dis(k,v)+1) {
            getans(v); break ;
        }
    }
}

int main() {
    n=rd(),m=rd();
    for(int i=1;i<=n;i++) se[bel[i]=rd()].eb(i);
    for(int i=1,u,v;i<n;i++) u=rd(),v=rd(),edg[u].eb(v),edg[v].eb(u);
    pre(1,0),init();
    for(int i=1;i<=m;i++) {
        int res=top[se[i][0]];
        for(int u:se[i]) if(top[u]^res) {printf("NE\n"); return 0;}
    }
    for(int i=1;i<=m;i++) {int k=top[se[i][0]]; dfs(k,0,0),dfz(fnd(k,0,sz[k]));}
    for(int i=1;i<=n;i++) cnt[i].resize(mxd[i]+1),tnc[i].resize(mxd[tp[i]]+1);
    for(int i=2;i<=n;i++) if(bel[i]^bel[ft[i]]) son[bel[ft[i]]].eb(bel[i]);
    solve(bel[1]);
    for(int u:se[bel[1]]) if(f[u]) {
        printf("DA\n"),getans(u);
        for(int i=1;i<=m;i++) printf("%d%c",ans[i],"\n "[i<m]);
        return 0;
    }
    printf("NE\n");
    return 0;
}