(15647) 로스팅하는 엠마도 바리스타입니다

https://www.acmicpc.net/problem/15647

단계별로 풀어보기

36단계(트리에서의 동적 계획법) 5번째

 

 

 

1번 농장을 루트 노드로 가정하자. 1번 농장을 기준으로 모든 농장까지의 거리를 구할 수 있고, 거리의 총합도 구할 수 있다.

x번 농장을 루트로 하는 서브트리의 노드 수를 cnt[x]라고 하자.

이때 x번 농장에서 y번 농장으로 기준이 바뀌었을 때 거리 총합의 변화량을 계산하면

cnt[y]개의 농장은 거리가 d만큼 줄어들고, n-cnt[y]개의 농장은 거리가 d만큼 늘어난다.

따라서 x번 농장 기준 거리 총합에 (n-2*cnt[y])*d를 더하면 y번 농장 기준 거리 총합이 된다.

처음에 1번 농장을 기준으로 한 거리 총합을 구했으므로, DFS를 돌리며 모든 농장 각각을 기준으로 한 거리 총합을 갱신해주면 된다.

 

 

 

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
typedef pair<int,int> p;

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
  
  int n;
  
  cin >> n;
  
  vector adj(n+1, vector<p>());
  vector<int> dist(n+1), cnt(n+1);
  vector<long> res(n+1);
  
  for(int i=1;i<n;i++){
    int u, v, d;
    
    cin >> u >> v >> d;
    
    adj[u].push_back({d,v});
    adj[v].push_back({d,u});
  }
  
  auto dfs= [&](auto self, int now, int nowdist) -> int {
    dist[now]= nowdist;
    cnt[now]= 1;
    
    for(p next:adj[now]){
      if(cnt[next.second])  continue;
      
      cnt[now]+= self(self, next.second, nowdist+next.first);
    }
    
    return cnt[now];
  };
  
  auto calc= [&](auto self, int now) -> void {
    for(p next:adj[now]){
      if(res[next.second])  continue;
      
      res[next.second]= res[now] + (n-2*cnt[next.second])*next.first;
      
      self(self, next.second);
    }
  };
  
  dfs(dfs, 1, 0);
  
  for(int i=1;i<=n;i++)  res[1]+= dist[i];
  
  calc(calc, 1);
  
  for(int i=1;i<=n;i++)  cout << res[i] << "\n";
  
  return 0;
}