Codeforce 1401 D. Maximum Distributed Tree 解析(思維、DFS、組合、貪心、DP)

今天我們來看看CF1401D
題目連結

題目
直接看原題比較清楚，略。

前言

這次比賽被第C.題搞到剩20分鐘可以寫D.這題，比賽時沒寫出來，比完了以後花了一個多小時Debug才De出來Orz。

想法

要注意到，題目中的distribution index實際上就是把每條路對於每個點對互相拜訪時，會被經過幾次，乘個權重，全部加起來。例如說$(u,v)$這條邊，$u$那一端接了一個節點個數為$cnt_u$的子樹，$v$那一端接了一個節點個數為$cnt_v$的子樹(code中我是用$dp[.]$來儲存)，那麼總共有$cnt_u\times cnt_v$個點對，在互相拜訪時會經過$(u,v)$這條邊，而distribution index就是把所有邊的拜訪次數乘以你分配的權重，然後全部加起來。
而對於每條邊，如果要找到兩個節點分別連接到多大的子樹(注意，這邊所說的子樹代表由當前點為根，且不會經過目前考慮的邊的子樹)，我們可以隨便選一個點$r$為根，做樹上DP維護每個點為根，往下的子樹(也就是以$r$為根來說，我們考慮的子樹不會往$r$接近)的節點個數。接著就是一些簡單的實作了。
要注意以下幾點:

1. 對於某個邊$(u,v)$，這個邊會被經過的次數為:$dp[v]\times(n-dp[v])$ (假設$u$比$v$淺，需要先在dfs時維護每個點的depth)
2. 最後把權重分配到邊上時，首先當然兩個數列都要先排序，接著要記得考慮$n-1>m$和$n-1\le m$的情況(我就是在這裡卡了一個多小時)

程式碼:

int _n=1e5+10;  
int t,n,m,p[_n],x,y,dp[_n],dep[_n];  
VI G[_n];  
PII e[_n];  
ll ecnt[_n];  
void dfs(int v,int fa,int d){  
  dep[v]=d;  
  if(SZ(G[v])==1 and G[v][0]==fa){dp[v]=1;return;}  
  int res=1;  
  rep(i,0,SZ(G[v]))if(fa!=G[v][i])dfs(G[v][i],v,d+1),res+=dp[G[v][i]];  
  dp[v]=res;  
}  
main(void) {ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);  
  cin>>t;while(t--){  
    cin>>n;rep(i,0,n)G[i].clear();  
    rep(i,0,n-1){cin>>x>>y;x--,y--;G[x].pb(y),G[y].pb(x);e[i]={x,y};}  
    cin>>m;rep(i,0,m)cin>>p[i]; sort(p,p+m);  
    dfs(0,-1,0); ll ans=0;rep(i,0,n-1){  
      int u=e[i].fi,v=e[i].se;  
      if(dep[u]>dep[v])swap(u,v);  
      ecnt[i]=1ll*dp[v]*(n-dp[v]);  
    } sort(ecnt,ecnt+n-1);  
    int len=max(0,n-1-m);  
    rep(i,0,len)ans+=ecnt[i],ans%=mod;  
    if(n-1>m)rep(i,len,len+m)ans+=1ll*p[i-len]*ecnt[i],ans%=mod;  
    if(m>=n-1){  
      rep(i,0,n-2)ans+=1ll*p[i]*ecnt[i],ans%=mod;  
      ll tmp=1;rep(i,n-2,m)tmp*=1ll*p[i],tmp%=mod;  
      ans+=1ll*tmp*(ecnt[n-2]%mod); ans%=mod;  
    }  
    cout<<ans<<'\n';  
  }  
  return 0;  
}  

標頭、模板請點Submission看
Submission