[백준] 4195. 친구 네트워크 - Python

[Gold II]

 

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

4195번: 친구 네트워크

 

풀이

Union-Find를 이용한, 서로소 집합(Disjoint Set) 응용 문제.

 

기존 Union-Find를 이용한 기본적인 문제와 크게 다르지 않으나,

네트워크에 속한 친구의 수를 따로 계산해줘야 한다. (rank와는 다름.)

 

네트워크에 속한 친구의 수를 기록하는 dictionary를 선언하고, (cnt)

union해 주는 과정에서, root가 되는 쪽에다 다른 쪽의 네트워크 친구 수를 더해버리면 된다.

이 과정에서, 기존 union 함수에서 예외처리 하지 않았던 x==y 일 경우를 예외처리해야 한다.

하지 않으면, 이미 친구 관계임에도 네트워크 친구 수가 다시 한 번 더 더해지는 오류를 범할 수 있다.

 

AC.

T = int(input())

root = {}
cnt = {}

def find(x):
    if root[x] == x:
        return x
    # 경로 압축
    root[x] = find(root[x])
    return root[x]

def union(x, y):
    x = find(x)
    y = find(y)
    
    if x == y:
        return
    
    if x > y:
        root[x] = y
        # y를 root로 설정했으니, y의 네트워크 수에 x의 네트워크 수를 더해서 저장.
        cnt[y] += cnt[x]
    else:
        root[y] = x
        # x를 root로 설정했으니, x의 네트워크 수에 y의 네트워크 수를 더해서 저장.
        cnt[x] += cnt[y]

for _ in range(T):
    F = int(input())
    root = {}   # Union-Find root
    cnt = {}    # 친구 네트워크에 있는 사람 수
    
    for _ in range(F):
        a, b = input().split()
        if a not in root:
            root[a] = a
            cnt[a] = 1
        if b not in root:
            root[b] = b
            cnt[b] = 1
        union(a, b)
        print(cnt[find(a)])