LGP10060 [SNTS 2024] 树上Voronoi图 学习笔记

LGP10060 [SNTS 2024] 树上Voronoi图 学习笔记

Luogu Link

题意简述

你有一棵 \(n\) 个结点的无根树。定义树上两点间的距离为其间简单路径的边数。

有 \(k\) 个关键点 \(a_1,a_2,\dots,a_k\)。我们记 \(f(u)\) 为距离 \(u\) 最近的关键点的关键点序号，即满足 \(\text{dis}(u,a_i)\) 尽量小的 \(i\)。特别地，若有多个 \(i\) 满足要求，则 \(f(u)\) 取其中最小的 \(i\)。

现在已知 \(F\)。问 \(A\) 有多少种可能方案。\(n\le 3\times 10^3\)。

做法解析

有一个事情是显然的：\(f(u)\) 相同的点应该连成一片（实际上 \(f(u)\) 的取值可以看作一种颜色）。

但是答案并不是所有连通块的大小乘起来就完了。因为可能会出现“\(f(u)=1\) 的点离 \(a_2\) 更近”此类情况。

我们怎么考察这个呢？有一个重要的性质：如果出现不合法状况，那么一定有不合法状况存在于颜色的交界处（否则这就与“同色点必然连成一片”矛盾）。

于是……是的，我们可以直接 \(O(n^2)\) 检查所有颜色两两间的交界处。还等什么，准备树形 \(\texttt{DP}\) 吧。

对于这道题来说，一个 \(\texttt{DP}\) 方案是这样的：先把相同颜色的点“缩到一起”（只是为了把颜色当成结点来考虑动规），然后定义 \(f(i,u)\) 为：第 \(i\) 种颜色取点 \(u\) 为关键点时，其（在缩点后的树上）自己和子树的选择方案总数。

我们每访问一个颜色 \(i\)，首先去把它的子问题做完，然后我们枚举：\(i\) 的关键点 \(u\)、相邻的颜色 \(j\)、\(j\) 的关键点 \(v\)。啊就是这样，做完了。

反正时间复杂度是 \(O(n^2)\) 的。这题 \(O(n^2)\) 就是对的。

代码实现

不算难也不算易。

#include <bits/stdc++.h>
using namespace std;
using namespace obasic;
using namespace omodint;
using mint=m998;
const int MaxN=3e3+5;
int N,K,X,Y,F[MaxN];bool flg;
int brd[MaxN][MaxN],vis[MaxN];
int dis[MaxN][MaxN];mint dp[MaxN][MaxN],ans;
vector<int> Tr1[MaxN],S[MaxN],Tr2[MaxN];
void addudge1(int u,int v){
    Tr1[u].push_back(v);
    Tr1[v].push_back(u);
}
void befinit(int n,int k){
    flg=0,ans=0;for(int i=1;i<=n;i++)Tr1[i].clear();
    for(int i=1;i<=k;i++)S[i].clear(),Tr2[i].clear(),vis[i]=0;
}
void dfs1(int u,int f,bool &flag){
    if(flag)return;
    if(F[u]!=F[f]){
        Tr2[F[f]].push_back(F[u]);
        Tr2[F[u]].push_back(F[f]);
        brd[F[u]][F[f]]=u,brd[F[f]][F[u]]=f;
        vis[F[u]]++;if(vis[F[u]]>1)flag=1;
    }
    S[F[u]].push_back(u);
    for(auto v : Tr1[u])if(v!=f)dfs1(v,u,flag);
}
void dfs2(int rt,int u,int f,int d){
    dis[rt][u]=d;
    for(int v : Tr1[u])if(v!=f)dfs2(rt,v,u,d+1);
}
bool check(int i,int u,int j,int v){
    if(i>j)swap(i,j),swap(u,v);
    int ib=brd[i][j],jb=brd[j][i];
    return (dis[ib][u]<=dis[ib][v])&&(dis[jb][v]<dis[jb][u]);
}
void dfs3(int i,int f){
    for(int j : Tr2[i])if(j!=f)dfs3(j,i);
    for(int u : S[i]){
        dp[i][u]=1;
        for(int j : Tr2[i]){
            if(j==f)continue;mint csum=0;
            for(int v : S[j])if(check(i,u,j,v))csum+=dp[j][v];
            dp[i][u]*=csum;
        }
    }
}
void mian(){
    readis(N,K);befinit(N,K);
    for(int i=1;i<N;i++)readis(X,Y),addudge1(X,Y);
    for(int i=1;i<=N;i++)readi(F[i]);
    dfs1(1,0,flg);if(flg){puts("0");return;}
    for(int i=1;i<=N;i++)dfs2(i,i,0,0);
    int rt=F[1];dfs3(rt,0);for(auto v : S[rt])ans+=dp[rt][v];
    writil(miti(ans));
}
int Tcn;
int main(){
    readi(Tcn);
    while(Tcn--)mian();
    return 0;
}