AT_agc034_e Complete Compress
这里是用 set 实现的换根 DP，时间复杂度 \(O(n\log n)\)。
记 \(siz_x,g_x,f_x\) 分别为 \(x\) 及其子树中有多少个关键点，所有关键点到 \(x\) 的距离和，将关键点尽可能两两向上合并后到 \(x\) 的距离和（我愿意理解为是将 \(g_x\) 缩到最小的值）。
当你钦定最后合并到 \(root\)，那么是否有解等价于 \([f_{root}=0]\)。
考虑朴素 DP，枚举所有根。有转移：
- \(siz_x=[a_x=1]+\sum\limits_{y\in son(x)} siz_y\)
- \(g_x=\sum\limits_{y\in son(x)} (g_y+siz_y)\)
- \(f_x=\begin{cases} f_y+siz_y-(g_x-(g_y+siz_y)) & \exists y\in son(x),f_y+siz_y > g_x-(g_y+siz_y) \\ g_x\bmod 2 & \forall y\in son(x),f_y+siz_y\leq g_x-(g_y+siz_y) \end{cases}\)
\(\exists y\in son(x),f_y+siz_y > g_x-(g_y+siz_y)\)，实际含义是如果有一个 \(y\)，它及其子树到 \(x\) 所需的最少步数比 \(x\) 其他子树所提供的最多步数还多就不能被减完，最优就是尽可能的多处理这个子树；不然就考虑每次选出所有剩余 \(f_y\) 中最大的两个，一起 \(-1\)，这样就相当于最后只会剩 \([0,1]\) 个 \(1\)，就等于 \(g_x\) 的奇偶性。
考虑接下来换根 DP，先考虑每个 \(x,y\in son(x),f_y+siz_y > g_x-(g_y+siz_y)\) 最多只会有 \(1\) 个，所以可以移项，判断 \(f_y+siz_y+g_y+siz_y\) 与 \(g_x\) 的大小关系。\(f_x\) 的第一种转移也可以变成 \(f_y+siz_y+g_y+siz_y-g_x\)。
所以 \(siz_x,g_x\) 正常换根，再对于每个点开一个 set，将所有的 \(f_y+siz_y+g_y+siz_y\) 放到 \(x\) 对应的 set 中，这样就可以换根了。
其实我们每次只用到了最大值，所以如果不用 set 也可以直接维护，就可以 \(O(n)\) 了，但我感觉很麻烦。
#include<stdio.h>
#include<iostream>
#include<algorithm>
#include<vector>
#include<set>
using namespace std;
const int MAXN=1e6+10;
const long long INF=1e16+10;
int n,siz[MAXN];char a[MAXN];
long long f[MAXN],g[MAXN],ans=INF;
vector <int> v[MAXN];set <long long> s[MAXN];
inline void clear()
{
for(int i=1;i<=n;++i) v[i].clear(),s[i].clear();
ans=INF;return ;
}
void get_f(int x)
{
if(s[x].empty()){f[x]=0;return ;}
if(*s[x].rbegin()<=g[x]) f[x]=(g[x]&1?1:0);
else f[x]=*s[x].rbegin()-g[x];return ;
}
void dfs(int x,int fa=0)
{
siz[x]=(a[x]=='1');g[x]=0;
for(int y:v[x])
{
if(y==fa) continue;dfs(y,x);
siz[x]+=siz[y],g[x]+=g[y]+siz[y];
s[x].insert(f[y]+siz[y]+g[y]+siz[y]);
}
get_f(x);return ;
}
void dp(int x,int fa=0)
{
if(!f[x]) ans=min(ans,g[x]/2);
for(int y:v[x])
{
if(y==fa) continue;
siz[x]-=siz[y],g[x]-=g[y]+siz[y];
s[x].erase(s[x].find(f[y]+siz[y]+g[y]+siz[y])),get_f(x);
siz[y]+=siz[x],g[y]+=g[x]+siz[x];
s[y].insert(f[x]+siz[x]+g[x]+siz[x]),get_f(y);
dp(y,x);
siz[y]-=siz[x],g[y]-=g[x]+siz[x];
s[y].erase(s[y].find(f[x]+siz[x]+g[x]+siz[x])),get_f(y);
siz[x]+=siz[y],g[x]+=g[y]+siz[y];
s[x].insert(f[y]+siz[y]+g[y]+siz[y]),get_f(x);
}
return ;
}
int main()
{
freopen("charlotte.in","r",stdin);
freopen("charlotte.out","w",stdout);
cin.tie(0),cout.tie(0);
ios::sync_with_stdio(0);
while(cin>>n)
{
for(int i=1;i<=n;++i) cin>>a[i];
for(int i=1,x,y;i<n;++i)
cin>>x>>y,v[x].push_back(y),v[y].push_back(x);
dfs(1),dp(1);
cout<<(ans==INF?-1:ans)<<'\n';clear();
}
return 0;
}
