백준 4195번 : 친구 네트워크

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

4195번: 친구 네트워크

 

처음 접해본 Union-Find 문제였다.

 

문제를 간단히 요악하면 두 사람의 이름이 입력으로 주어질 때 각 사람이 가진 친구들을 합한 수를 출력하는 문제이다. 예를 들어 A와 B가 입력으로 주어지면 그때부터 A와 B는 친구 관계를 맺게되고, A와 B가 친구가 되는 순간 A의 친구는 B의 친구가 되고, B의 친구는 A의 친구가 된다. 그 상태에서 A와 B가 가진 친구들의 수를 구해야 하는 것이다.

 

이름이 문자열로 주어지기 때문에 사람의 이름을 Key로 하고 그 사람이 가진 친구들의 이름을 담을 수 있는 Set를 Value로 하는 사전을 이용하려고 했다. 집합을 사용하려고 한 이유는 입력으로 주어진 두 사람과 모두 친구인 사람이 있을 수 있어서 합쳐졌을 때 중복을 제거하기 위해서였다. 문제는 입력으로 주어진 두 사람의 이름을 A와 B라고 했을때 A와 친구인 사람이 C,D고 B와 친구인 사람이 E,F라고 한다면, A가 가진 친구 수를 계산하기 위해 C의 친구, D의 친구를 모두 확인해야했다. N이 증가할수록 확인하는데 시간이 오래 걸릴 것 같아서 A와 직접적으로 친구를 맺은 사람뿐만 아니라 간접적으로 친구가 된 사람들의 이름도 집합에 넣어가며 계산했다. 결과는 시간 초과를 메모리 초과로 치환한 꼴밖에 되지 않았다.

 

결국 Union-Find 알고리즘을 사용해야한다는 힌트를 얻고 다시 문제를 풀었다. Union-Find는 서로소인 두 집합을 합칠 때 사용하는데, 단순히 참조하는 루트 노드를 바꿈으로써 집합을 합칠 수 있다. 그러니까 A의 친구가 C,D 라는 것은 C,D의 루트노드가 A로 설정돼있다는 의미이다. 이렇게 하면 C,D와 친구인 사람까지 고려하지 않아도 된다. 왜냐하면 C,D의 루트 노드가 A이기 때문에 C,D를 루트로 하는 다른 노드들 역시 루트를 따라 올라가면 A와 친구임이 증명되기 때문이다.

 

예를들어 하나의 조직을 상상해보자. A는 보스이고 그를 따르는 C, D는 조직원이다. 이때 C와 D를 따르는 더 말단인 조직원이 있을 수 있다. 그렇지만 결국 이 말단 조직원들조차 자신들이 따르는 사람이 A를 따르기 때문에 전체적으로 보면 같은 조직(친구)에 소속되있다는 것은 분명하다. 그런데 이 조직이 다른 조직과 싸워서 진 쪽이 이긴 쪽의 아래로 들어간다고 생각해보자. 이 때 싸움에서 진 조직원들은 자신들의 소속을 일일히 변경할 필요없이 싸움에서 진쪽의 보스가 이긴쪽의 보스를 따르면 인수인계는 끝나는 것이다. 인수인계 내용은 보스가 갖고있는 조직원들의 수이다. 결국 입력으로 주어진 두 사람의 소속이 다를 때 (조직의 보스가 다를 때) 조직을 하나로 통합하고 통합된 조직의 조직원 수를 출력하는 것이다. ( parent 함수에서 재귀를 통해 조직원의 보스를 찾을때 수행 시간을 단축할 수 있게 작성했다. )

 

import sys
input = sys.stdin.readline

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

def union(x, y):
    px = parent(x)
    py = parent(y)
    if px != py:
        p[py] = px
        c[px] += c[py]

for _ in range(int(input())):
    n = int(input())
    p, c = {}, {}
    for _ in range(n):
        f1, f2 = input().split()
        if f1 not in p.keys():
            p[f1] = f1
            c[f1] = 1
        if f2 not in p.keys():
            p[f2] = f2
            c[f2] = 1
        union(f1, f2)
        print(c[parent(f1)])