[백준] 1289 트리의 가중치

문제 링크 : www.acmicpc.net/problem/1289

1289번: 트리의 가중치

문제 유형

Tree DP, Math

 

어떤 한 서브트리에서 나올 수 있는 모든 경로를 구해 보자.

우선 root를 경유하지 않는 경로, 즉 root노드가 경로의 한쪽 끝인 경우이다. 이는 일반적인 재귀를 통해 모두 구할 수 있다.

문제는 root를 경유하는 경로이다.

이를 구하기 위해서는 root로부터 이어진 각 자식들의 서브트리 가중치를 서로 이어줘야 하기 때문이다.

그림의 (왼쪽 자식 노드를 root로 하는 서브트리의 가중치 * root와 왼쪽 자식 노드 사이의 간선 가중치)를 L, 오른쪽의 동일한 값을 R이라 하자.

실제로는 L = (l1 + l2 + l3 + ...)의 형태이다. 서브트리에서 재귀적으로 모든 경로의 가중치를 다 더한 값을 L과 R로 정의한 것이다. 이 때, 위 트리의 가중치는 L x R이 된다.

하지만 전체 트리 root의 자식 노드 수가 많아지면 난감해진다. 모든 자식 노드들을 한 쌍씩 연결하는 Naive한 방법은 O(N^2)이므로, 시간 제한을 만족할 수 없기 때문이다.

따라서, 다음과 같은 과정을 통해 모든 쌍 연결을 O(N)에 수행해야 한다.

각 자식 노드들을 root로 하는 (서브트리 가중치 * 간선의 가중치) 값을 A, B, C, ... 라 할 때,

① 처음 A, B를 잇는다. 현재까지 구한 쌍은 (AB) 이다. 이제 A + B의 값을 저장해 둔다.

② 저장해둔 (A + B)에 C를 곱한다. 현재까지 구한 쌍은 (AB, AC, BC) 이다. 이제 저장해 둔 값에 C를 더한다.

③ 저장해둔 (A + B + C)에 D를 곱한다. 현재까지 구한 쌍은 (AB, AC, BC, AD, BD, CD) 이다. 저장해 둔 값에 D를 더한다.

....

각 과정들에서 곱해서 구한 값들을 모두 더하면 구하는 정답이 된다.

 

[코드]

#define _CRT_SECURE_NO_WARNINGS
#include <iostream>
#include <vector>
#include <algorithm>
 
#define DIV 1000000007
 
using namespace std;
typedef pair<int, int> p;
typedef long long ll;
 
int n;
vector<p> adj[100001];
ll answer;
 
ll get_path(int node, int pnode) {
    ll mult = 1;
    for (int i = 0; i < adj[node].size(); i++) {
        int child = adj[node][i].first;
        int dist = adj[node][i].second;
        if (child == pnode) continue;
 
        ll res = (get_path(child, node) * dist) % DIV;
        answer += (res * mult) % DIV;
        answer %= DIV;
        mult = (mult + res) % DIV;
    }
    return mult;
}
 
void init() {
    int a, b, w;
    cin >> n;
 
    for (int i = 0; i < n - 1; i++) {
        cin >> a >> b >> w;
        adj[a].push_back({ b, w });
        adj[b].push_back({ a, w });
    }
}
 
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(NULL);
 
    init();
    get_path(1, 0);
    printf("%lld\n", answer);
 
    return 0;
}