1289번: 트리의 가중치

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


$O(n)$

한 노드의 자식노드 a1,a2, ... 이 있다 하자.
모든 경우를 따지기 위해선
1. 자식노드 - 현재 노드 - 자식 노드
2. 자식노드 - 현재 노드
이 두 경우를 따져서 가중치를 구해 더하면 된다.

각 자식들로 끝나는 가중치를 bi 라 하자.
2번 케이스는 bi 전체 합을 이용해 쉽게 구할 수 있다.
1번 케이스는 bi-1 까지의 누적 합 곱하기 bi를 누적함으로써 구할 수 있다.

#include<stdio.h>
#include<vector>
#define mod 1000000007
using namespace std;
typedef long long ll;
const int MAX_N = 1e5;
vector<pair<int, int> > adj[MAX_N + 1];
int n, res;
int dfs(int h, int p) {
    int s = 1, t;
    for (auto it : adj[h]) {
        if (it.first == p) continue;
        t = (ll)dfs(it.first, h)*it.second%mod;
        res = (res + (ll)t*s) % mod;
        s = (s + t) % mod;
    }
    return s;
}
int main() {
    scanf("%d", &n);
    for (int i = 0, a, b, c; i < n - 1; i++) {
        scanf("%d %d %d", &a, &b, &c);
        adj[a].push_back({ b,c });
        adj[b].push_back({ a,c });
    }
    dfs(1, -1);
    printf("%d", res);
    return 0;
}