[백준][Python] 1289 트리의 가중치

트리dp.. 아직 잘 익지는 않는다. 좀 더 연습해야할듯.

 

import sys
input =sys.stdin.readline
sys.setrecursionlimit(10**6)

# dp[u]는 노드 u가 루트인 subtree에서 u부터 다른 모든 노드 까지 가는 모든 경로들의 곱의 합.
# 즉 (dp[v] * c) 들의 합. 그 값을 리스트 p에 저장해뒀다가 모든 값들에 대해 (dp[u] - dp[v]*c)*(c*dp[v])들의 합을 구해준다.
# 그 후 중복을 제거하기 위해 나누기 2를 해줘야 하나 MOD의 반인 500000004를 곱하고 MOD로 나눔으로써 2로 나누는 것과 같은 효과를 얻을 수 있다.
def dfs(u):
    global ans
    visited[u] = True
    p = []
    for i in range(len(edges[u])):
        v = edges[u][i]
        c = costs[u][i]
        if visited[v]:
            continue

        dfs(v)
        dp[u] += (dp[v] * c) % MOD
        dp[u] %= MOD

        p.append((dp[v] * c) % MOD)
    ans += dp[u]
    ans %= MOD

    sum_v = 0
    for v in p:
        sum_v += ((dp[u] - v + MOD) % MOD * v) % MOD
        sum_v %= MOD
    
    sum_v *= 500000004
    sum_v %= MOD
    
    ans += sum_v
    ans %= MOD                

    dp[u] += 1
    dp[u] %= MOD


MOD = 1000000007
N = int(input())
ans = 0
visited = {i: False for i in range(1,N+1)}
dp = [0 for _ in range(N+1)]
edges = {i: [] for i in range(1,N+1)}
costs = {i: [] for i in range(1,N+1)}
for _ in range(N-1):
    a, b, c = map(int, input().split())
    edges[a].append(b)
    edges[b].append(a)
    costs[a].append(c)
    costs[b].append(c)

dfs(1)
print(ans)