[BOJ/PYTHON] 4803. 트리

백준 문제 링크: https://www.acmicpc.net/problem/4803

문제 요약

트리는 사이클이 없는 연결 요소.

그래프 주어졌을 때 트리 개수 세기

핵심 아이디어

트리의 개수는 서로소 집합을 찾는 union-find로 풀 수 있다.

입력받은 간선의 각 부모 노드가 다르면 트리를 합쳐주고, 부모 노드가 같으면 사이클이 발생하는 그래프이기 때문에 사이클 리스트에 해당 노드를 추가해준다.

그 다음 모든 노드의 부모 노드를 갱신해주는 과정이 필요하다.

중복 제거한 부모 노드 리스트에서 해당 부모 노드가 사이클이 발생하지 않는다면 트리 개수를 추가해준다.

풀이

# gold4-4803. 트리

import sys
input = sys.stdin.readline

def find_parent(parent, x):
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])
    return parent[x]

def union_parent(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b: parent[b] = a
    else: parent[a] = b

case = 1
while True:
    n, m = map(int, input().split())
    if n == 0 and m == 0: break
    parent = [i for i in range(n + 1)]
    cycles = []
    cycle_parent = set() # 사이클이 발생한 노드의 부모 노드
    tree = 0
    for _ in range(m):
        a, b = map(int, input().split()) # 간선
        if find_parent(parent, a) != find_parent(parent, b): # 부모가 다르면 트리 합쳐줌
            union_parent(parent, a, b)
        else: cycles.append(a) # 부모가 같으면 사이클 리스트에 추가
    
    for i in range(n+1):
        find_parent(parent, i)

    for cycle in cycles:
        cycle_parent.add(parent[cycle])
    
    for i in set(parent): # 중복 제거한 부모 노드 리스트에서
        # 부모 노드가 사이클이 없으면 
        # (i != 0은 초기값 세팅이 0이라서 넣어줌. 0번 인덱스의 값인 0을 제외)
        if i != 0 and i not in cycle_parent: 
            tree += 1 # 트리 개수 추가

    if tree == 0: print(f'Case {case}: No trees.')
    elif tree == 1: print(f'Case {case}: There is one tree.')
    else: print(f'Case {case}: A forest of {tree} trees.')

    case += 1

Trouble Shooting

입력받을 때의 if find_parent(parent, a) != find_parent(parent, b) 로직 때문에 find_parent를 한 번씩 돈다고 생각해 find_parent 갱신을 안 해줘도 되는줄 알았는데 아니었다.

실행 시간

메모리 34292KB, 시간 260ms (python3)