백준1289 트리의 가중치

작성일 | In

문제 링크

  • http://icpc.me/1289

문제 출처

  • 2008 COI Final Exam #1 2번

사용 알고리즘

  • Tree DP

시간복잡도

  • $O(N)$

풀이

서브트리 안에서 발생하는 가중치의 합은 아래에서 다 처리하고, 현재 노드를 지나는 경로의 가중치들만 생각합시다.
이진 트리라고 생각해도 상관 없습니다.

현재 노드의 왼쪽에 있는 자손들의 가중치를 $(a_1, a_2, a_3, \ldots)$, 오른쪽에 있는 자손들의 가중치를 $(b_1, b_2, b_3, \ldots)$라고 하면 현재 노드를 지나는 경로들의 가중치의 합은 $(a_1+a_2+a_3+\ldots) \times (b_1+b_2+b_3+\ldots)$가 됩니다.
이 식을 이용해 열심히 구현해주면 됩니다.

쓰고보니 Tree DP가 아닌 것 같기도 하네요.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
#include <bits/stdc++.h>
#define x first
#define y second
#define all(v) v.begin(), v.end()
#define compress(v) sort(all(v)), v.erase(unique(all(v)), v.end())
using namespace std;

typedef long long ll;
typedef pair<ll, ll> p;

const ll mod = 1e9+7;

int n;
vector<p> g[101010];
ll ans;

ll f(int v, int b){
    ll now = 1;
    for(auto i : g[v]) if(i.x != b){
        ll t = f(i.x, v) * i.y % mod;
        ans += now * t % mod; ans %= mod;
        now += t; now %= mod;
    }
    return now;
}

int main(){
    ios_base::sync_with_stdio(false); cin.tie(nullptr);
    cin >> n;
    for(int i=1; i<n; i++){
        int s, e, x; cin >> s >> e >> x;
        g[s].emplace_back(e, x);
        g[e].emplace_back(s, x);
    }
    f(1, 0);
    cout << ans;
}