[백준] 17073 나무 위의 빗물 python

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

17073번: 나무 위의 빗물

트리 문제이다.

입력으로 주어지는 값은 양방향 그래프이고 인접리스트를 사용해서 노드끼리 연결 시킨다.

그리고 트리를 순회할때 parent 노드와 다른 번호의 노드만 순회를 하면 leaf 노드를 찾을수 있다.

leaf node를 찾으면 1을 리턴해주고 이 값들을 모두 더해주면 leaf 노드의 개수를 찾을수 있다.

import sys; input = sys.stdin.readline
sys.setrecursionlimit(10**9)
def order(cur_node, parent, graph):
    if cur_node != 1 and len(graph[cur_node]) == 1:
        # leaf node has a connection with parent node
        return 1

    ret = 0
    for son in graph[cur_node]:
        # check child not equal parent, because graph is not tree
        if son == parent: continue
        ret += order(son, cur_node, graph)

    return ret

def main():
    # 0 input
    N, W = map(int, input().split())
    g = [[] for _ in range(N + 1)]
    for _ in range(N - 1):
        # bidirectional graph
        # if using tree, check child not equal parent
        u, v = map(int, input().split())
        g[u].append(v)
        g[v].append(u)

    # 1 find th num of leaves tree traversal
    num_leaf = order(1, -1, g)

    # 2 output
    print(W/num_leaf)


if __name__ == '__main__':
    main()

 

이문제는 리프노드의 개수만 알면 되기 때문에 연결된 개수만 확인하면된다.

그래프를 직접 연결하지 않고 연결 개수만 증가시키고, 연결이 딱 하나만 있는 노드의 개수만 세면 된다.

하지만 연결이 하나 있는 노드에 루트 노드가 포함되는 경우가 가끔 생긴다.

순회하기전 루트 노드에 아무값이나 더해도 되고, 그래프를 연결시킨뒤에 root 노드만 1이 아닌 값으로 변경해도 된다.

import sys; input = sys.stdin.readline

def main():
    # 0 input
    N, W = map(int, input().split())
    g = [0 for _ in range(N + 1)]
    # if finding nodes have one connection, we find leaves. sometimes include root.
    # if a tree has one root and one leaf, each node has one connection.
    for _ in range(N - 1):
        # calculate the num of connection
        u, v = map(int, input().split())
        g[u] += 1
        g[v] += 1
    g[1] = 0 # remove root node
    # 1 find the num of leaves tree traversal
    num_leaf = sum([n for n in g if n == 1])

    # 2 output
    print(W/num_leaf)


if __name__ == '__main__':
    main()