题解：AT_arc108_f [ARC108F] Paint Tree

· · 题解

\text{Link}

应该比较自然的做法。

::::info[Hint1] 考虑树的直径有什么性质。 ::::

::::info[Hint2] 当直径端点异色的时候，答案是不是一定会选某一个直径端点。 ::::

::::info[Hint3] 考虑对两个端点分别统计贡献。 ::::

思路

如果某个方案直径两端点颜色相同，这种方案的贡献一定是直径长度，一共有 2^{n-1} 种这样的方案。

否则直径端点一黑一白。

发现最远同色点对一定可以取某个直径端点，不然可以反证直径长度会变长。

令直径两端点为 s 和 t，以 s 为例，钦定 s 与 u 是最远同色点对。

设 s 为白色，把所有点按照到 s 的距离从大到小排序，排在 u 前面的一定要是黑色的，u 一定要是白色的。到 t 距离大于 u 到 s 距离的也要是白色的，没被钦定颜色的点可以随便选，要特判一下一个点又被钦定了白色又被钦定了黑色，此时无解。

其它情况都是对称的，对于距离相等的情况我们钦定一下某个是更大的即可。

复杂度瓶颈是排序，桶排一下即可做到 O(n)。

可能有点 corner，代码里写的很清楚了。

Code

#include<bits/stdc++.h>
#define pb emplace_back
using namespace std;
const int N=2e5+10,mod=1e9+7;
int n,ma,to,s,t,disS[N],disT[N],p[N],p2[N],sum[N],cnt[N];
vector<int>e[N];
inline void dfs(int u,int fa,int dis){
    if(dis>ma){
      ma=dis;
      to=u;
    }
    for(int v:e[u]){
      if(v==fa)continue;
      dfs(v,u,dis+1);
    }
}
inline void dfs2(int u,int fa,int dis[]){
    for(int v:e[u]){
      if(v==fa)continue;
      dis[v]=dis[u]+1;
      dfs2(v,u,dis);
    }
}
signed main(){
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    cin>>n;
    p2[0]=1;
    for(int i=1;i<=n;++i)p2[i]=2ull*p2[i-1]%mod;
    int u,v;
    for(int i=1;i<n;++i){
      cin>>u>>v;
      e[u].pb(v);
      e[v].pb(u);
    }
    dfs(1,0,0);
    s=to;
    ma=to=0;
    dfs(s,0,0);
    t=to;
    dfs2(s,0,disS);
    dfs2(t,0,disT);
    int ans=1ull*p2[n-1]*ma%mod,val=0;
    for(int i=1;i<=n;++i){
      if(i==s||i==t)continue;
      ++sum[disT[i]];
    }
    for(int i=n;i>=0;--i)sum[i]+=sum[i+1];
    for(int i=1;i<=n;++i){
      if(i==s||i==t)continue;
      ++cnt[disS[i]];
    }
    for(int i=n;i>=0;--i)cnt[i]+=cnt[i+1];
    for(int i=1;i<=n;++i){
      if(i==s||i==t)continue;
      p[cnt[disS[i]]--]=i;
    }
    for(int i=1;i<=n-2;++i){
      int u=p[i];
      if(val>disS[u])break;
      ans=(1ull*disS[u]*p2[n-1-i-sum[disS[u]+1]+(disT[u]>disS[u])]+ans)%mod;
      val=max(val,disT[u]);
    }
    val=0;
    fill(sum,sum+1+n,0);
    for(int i=1;i<=n;++i){
      if(i==s||i==t)continue;
      ++sum[disS[i]];
    }
    for(int i=n;i>=0;--i)sum[i]+=sum[i+1];
    fill(cnt,cnt+1+n,0);
    for(int i=1;i<=n;++i){
      if(i==s||i==t)continue;
      ++cnt[disT[i]];
    }
    for(int i=n;i>=0;--i)cnt[i]+=cnt[i+1];
    for(int i=1;i<=n;++i){
      if(i==s||i==t)continue;
      p[cnt[disT[i]]--]=i;
    }
    for(int i=1;i<=n-2;++i){
      int u=p[i];
      if(val>=disT[u])break;
      ans=(1ull*disT[u]*p2[n-1-i-sum[disT[u]]+(disS[u]>=disT[u])]+ans)%mod;
      val=max(val,disS[u]);
    }
    cout<<ans<<"\n";
    return 0;
}