[백준] No.1289 - 트리의 가중치 (C++, 트리)

문제

 

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

1289번: 트리의 가중치

 

 

풀이

 

solved.ac 난이도: Platium 3

 

트리에서의 다이나믹 프로그래밍으로 분류되었지만 풀어보니 본인의 풀이에는 다이나믹이 사용될 부분은 없었다.

 

트리의 모든 경로를 하나 하나 계산하기에는 특정 노드에서 모든 노드 하나하나의 경로 수를 계산해보면

N(특정 노드)*(N-1)(선택한 노드 이외의 모든 노드)의 경로 수가 존재한다.

이때 각 노드의 경로에 중첩되는 부분을 다이나믹 프로그래밍으로 해결하는 풀이도 있는 듯 하지만 다르게 풀어보자.

 

now를 현재 노드라 하고 now를 root로 하는 서브 트리의 트리 가중치를 구해보자.

 

now의 자식 노드(1~i)들 각각을 root로 하는 서브트리들을 생각해보자.

각각의 서브트리들에서, 모든 하위 노드들로 부터 root로의 모든 경로 가중치의 합을 T1, T2, ... , Ti라고 하자.

하위 노드들에서 시작해 서브 트리의 root로 도달하는 모든 경로로들을 now로 도달하도록 변경하고 싶다면 T1~Ti를 이용하면 된다.

now의 각각의 자식 노드들로의 가중치는 r1, r2, ..., ri라고 하면 

(i번째 자식노드)와 (i번째 자식노드의 모든 하위 노드)에서 now로의 경로 가중치의 합은 (1 + Ti )* ri 이다.

 

예를 들어 아래의 트리를 생각해보자.

 

 

now 노드의 자식 노드는 1,2,3 그리고 노드 3의 자식노드 4,5가 존재한다.

해당 트리에서 존재하는 모든 하위 노드에서 루트로의 경로는 위와 같이 7개가 존재한다.

 

노드 3의 서브 트리만을 생각해보자.

해당 서브 트리에서는 리프 노드에서 루트인 3으로의 경로 r4와 r5만이 존재한다.

따라서 서브트리의 루트로의 모든 경로 가중치 T3는 r4 + r5이다.

 

이때 이 서브트리의 노드들에서 now 노드로의 경로 가중치는 3가지이다.

1. 서브트리의 루트인 3에서 now로의 경로 가중치 r3 

2. 4에서 now로의 경로 가중치 r3* r4

3. 5에서 now로의 경로 가중치 r3* r5

 

즉, (1+ T3) * r3 ( = sub_tree_weight 3이라고 하자,)가 성립한다.

 

이제 now를 루트로한 모든 하위 노드에서의 경로 가중치를 구할 수 있다.

그렇다면 now의 트리와 각각의 서브트리의 트리 가중치는 어떻게 구할까?

now의 트리 가중치는 구해놓은 now로의 경로 가중치를 통해 구할 수 있다.

 

now를 거쳐 서브 트리1의 모든 노드들과 서브 트리 2의 모든 노드들 이루는 경로들의 가중치는

sub_tree_weight1(= (1+T1)* r1) * sub_tree_weight2(= (1+T2)* r2) 로 구할 수 있다.

 

sub_tree_weight 1 (= (1+Ti)* ri) 는

서브트리 1의 모든 하위 노드에서의 now까지의 경로들의 가중치의 합이 존재한다.

각각을 r1, r2, r3...라고 하자.

sub_tree_weight 2 또한

서브트리 2의 모든 하위 노드에서의 now까지의 경로들의 가중치의 합이 존재한다.

각각의 k1, k2, k3...라고 하자.

 

이때 now를 경유하여 각 서브 트리들이 이루는 모든 경로의 가중치들은 결국 서로를 곱해주므로 구할 수 있다.

(r1+ r2 + r3 ....) * (k1 + k2 + k3 ...) = sub_tree_weight 1 * sub_tree_weight 2

 

따라서 구해놓은  sub_tree_weight 1~i들에서 두 개씩 짝을 이루는 모두 경우를 곱해주면 된다.

만약 단순히 2중 for문으로는 i*(i-1)/2 만큼의 계산이 필요하니 구간합 알고리즘으로 빠르게 계산해주자.

 

 

for(int i=1; i<psum.size() ; ++i)
  psum[i] = (psum[i]+ psum[i-1]);
 
for(int i=2; i<psum.size() ; ++i)
  total_weight += ((psum[i-1] - psum[i-2]) * ((psum[psum.size()-1] - psum[i-1])%MOD)) %MOD;

 

자.

now를 루트로 하는 서브트리에서의 모든 경로의 가중치

즉, now의 트리 가중치는 

모든 하위 노드들에서 now로의 경로 가중치인  sub_tree_weight들과

위에서 psum(구간합)으로 구한 now를 경유하여 서브 트리끼리 이루는 모든 경로들이다.

 

이를 트리 전체의 트리 가중치에 더해주고 now의 모든 sub_tree_weight들의 합을 반환함으로써

node 서브트리의 트리 가중치를 total_weight에 더해주고 하위 노드에서 node로의 모든 경로 가중치를 반환하는 함수 

int TreeWeight(int node)를 구현할 수 있다.

 

 

int TreeWeight(int node){
    vector<long long> psum;
    psum.push_back(0);
    
    long long sub_tree_weight;
    for(int i=0; i<children.size() ; ++i){
      if(children[i].first != parent){
        sub_tree_weight = tree[children[i].first].TreeWeight(children[i].first);
        sub_tree_weight = (sub_tree_weight +1) * children[i].second;
        sub_tree_weight %= MOD;
        total_weight += sub_tree_weight;
        psum.push_back( sub_tree_weight );
      }
    }
    total_weight %= MOD;
 
    for(int i=1; i<psum.size() ; ++i)
      psum[i] = (psum[i]+ psum[i-1]);
 
    for(int i=2; i<psum.size() ; ++i)
      total_weight += ((psum[i-1] - psum[i-2]) * ((psum[psum.size()-1] - psum[i-1])%MOD)) %MOD;
      
    total_weight %= MOD;
 
    return psum[psum.size()-1]%MOD;
  }
 
};

 

 

MOD 나머지 연산은 적절히 추가해 주자.

문제의 분류를 보고 틀린 풀이를 한 걸까 의심했는데 좀 더 자신의 풀이에 자신감을 가져도 될 때가 온 것 같다.

 

 

 

코드