[골드 5] 백준 17072 - 나무 위의 빗물 (파이썬)

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

풀이

이 문제는 비의 양 W가 루트 노드에 고이며, 나무는 N개의 정점으로 이루어진 트리이고, 루트 노드는 1번이다.

비는 모든 리프 노드까지 동일하게 분배되며, 각 리프 노드까지 고인 빗물의 양을 구하는 문제이다.

결국 이 문제를 해결하는 방법은 다음과 같다.

1. 트리는 사이클이 없는 연결 그래프이므로, DFS로 리프 노드를 쉽게 찾을 수 있다.
2. 루트에서 DFS를 돌려 리프 노드의 개수를 센다.
3. 전체 물이의 양 W를 리프 노드 개수로 나눈다.
   1. 리프노드마다 동일한 양의 물이 고이기 때문이다.

먼저 정점 수, 물의 양을 입력받는 인접 리스트를 생성한다.

N, W = map(int, input().split())
adj = [[] for _ in range(N)]

N은 트리의 정점 개수, W는 전체 빗물의 양, adj는 트리의 구조를 표현하는 인접 리스트를 의미한다.

for _ in range(N - 1):
    u, v = map(int, input().split())
    u -= 1
    v -= 1
    adj[u].append(v)
    adj[v].append(u)

트리는 N - 1개의 간선을 가지며, 입력된 정점 번호는 1부터 시작하므로 0부터 시작하도록 0-Index based 형식으로 처리한다

무방향 그래프이므로 adj[u], adj[v] 양쪽에 추가한다.

방문 처리 배열 및 리프 노드 수를 초기화 한다.

visit = [False] * N
count = 0

visit의 경우 DFS 중복 방문 방지용 배열이며, count는 리프 노드의 개수를 셀 변수를 나타낸다.

다음은 DFS 함수를 정의한다.

def dfs(u):
    visit[u] = True
    is_leaf = True

    for v in adj[u]:
        if not visit[v]:
            dfs(v)
            is_leaf = False
    
    if is_leaf:
        global count
        count += 1

u는 현재 방문 중인 노드를 의미하며 is_leaf는 자식 노드가 있는지 확인한다.

자식이 있으면 리프가 아니므로 is_leaf = False로 리프 노드인 경우 count를 증가시킨다.

결국 전체적인 흐름은 다음과 같다.

1. 트리 구조를 입력받고 인접 리스트를 구성한다.
2. DFS로 루트부터 탐색하면서 리프 노드 개수를 파악한다.
3. 물의 양 W를 리프 노드 수 count로 나눠서 정답을 출력한다.

코드

# 백준 17073 - 나무 위의 빗물
# 분류 : DFS + 수학

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

N, W = map(int, input().split())
adj = [[] for _ in range(N)]

for _ in range(N - 1) :
    u, v = map(int, input().split())
    u -= 1
    v -= 1

    adj[u].append(v)
    adj[v].append(u)

visit = [False] * N
count = 0

def dfs(u) :
    visit[u] = True
    is_leaf = True

    for v in adj[u] :
        if not visit[v] :
            dfs(v)
            is_leaf = False
    
    if is_leaf :
        global count
        count += 1

dfs(0)
print(W / count)