문제 링크: https://www.acmicpc.net/problem/4195

분리 집합을 활용해 풀 수 있는 문제입니다. 단, 두 문자열 간의 연결 관계가 주어지므로 맵을 활용해야 합니다.

풀이

​분리 집합 기반의 Union-Find 알고리즘

​분리 집합은 흔히 다익스트라 구현에서 사이클 여부를 판단하기 위해 사용하는 자료구조로, 여러 원소들을 단일 집합으로 합치고, 임의의 두 원소가 서로 같은 집합에 있는지 판단할 수 있습니다.

이 문제도 분리 집합을 활용해, 주어진 친구 관계에 따라 두 친구(원소)를 같은 친구 네트워크(집합)에 합치고, 같은 집합에 속하는 원소의 수를 세면 됩니다.

구현

평소의 분리 집합은 배열을 기반으로 구현하지만, 이 문제는 원소가 문자열로 주어지므로 맵을 사용해야 합니다.

input = open(0).readline

names = {} # 각 사람의 이름별로, 자신이 속한 집합의 루트가 되는 원소를 저장할 맵
friend_count = {} # 각 사람의 이름별로, 자신이 루트가 되는 집합에 속한 원소의 수를 저장할 맵

def init_name(name):
    """names 및 friend_count 맵에 새로운 사람의 정보를 추가한다."""
    names[name] = name
    friend_count[name] = 1

def find_root(name):
    """name이 속한 집합의 루트 원소를 반환한다."""
    path = []
    while name != names[name]: # 루트 원소를 찾는다.
        path.append(name)
        name = names[name]
    for p in path: # names 맵의 루트 정보를 갱신한다. (탐색 깊이 줄이기)
        names[p] = name
    
    return name

def union(a, b):
    """두 원소 a, b가 속한 두 집합을 병합한다."""
    a_root = find_root(a)
    b_root = find_root(b)

    if a_root == b_root: # 이미 같은 집합의 원소인 경우, 병합하지 않는다.
        return False

    names[b_root] = a_root
    friend_count[a_root] += friend_count[b_root]
    return True

전체 코드

input = open(0).readline

names = {}
friend_count = {}

def init_name(name):
    names[name] = name
    friend_count[name] = 1

def find_root(name):
    path = []
    while name != names[name]:
        path.append(name)
        name = names[name]
    for p in path:
        names[p] = name
    
    return name

def union(a, b):
    a_root = find_root(a)
    b_root = find_root(b)

    if a_root == b_root:
        return False

    names[b_root] = a_root
    friend_count[a_root] += friend_count[b_root]
    return True

for _ in range(int(input())):
    N = int(input())
    names.clear()
    friend_count.clear()
    for _ in range(N):
        name1, name2 = input().rstrip().split()

        # 초기화
        if names.get(name1, None) is None:
            init_name(name1)
        if names.get(name2, None) is None:
            init_name(name2)
        
        # 병합
        union(name1, name2)
        print(friend_count[find_root(name1)])