백준 4195번 친구 네트워크 | python | union-find

문제

링크

풀이

보통 집합의 원소가 숫자로 되어 있는 경우가 많은데 여기선 문자열이 들어오기 때문에 딕셔너리를 사용해준다 <- 여기까진 스스로도 생각해내었음

근데 한 줄씩 입력될 때마다(친구 관계가 생길 때마다) 두 사람의 친구 네트워크에 몇 명이 있는지 세어주는 걸 어떻게 해야할 지 고민하다가

같은 부모를 갖는 딕셔너리의 개수를 count해주려 했는데 왜인지 오류가 나서 구글링의 힘을 빌렸다.

일단 기존 union 함수와 해당 문제에서의 union 함수의 차이는
부모 대소 비교를 할 수 없으므로 (숫자가 아닌 문자열이니까)
무조건 먼저 입력된 값을 상위 부모라고 생각하고 바꿔주는 것이다.

더불어 친구 관계도 먼저 입력된 값을 기준으로 탐색 횟수를 계속 더해주면 된다. 아마 아래 코드를 보면 오히려 이해가 쉬울 것이다.

import sys
sys.setrecursionlimit(10 ** 6)


# a가 속한 집합과 b가 속한 집합 합치기
def union(x, y):
    # 각 친구의 부모 노드를 찾는다.
    x = find(x)
    y = find(y)

    # 부모 노드가 같으면 리턴
    if x == y:
        return

    else:
        # 부모 노드가 다르면 y에 부모 노드를 x에 부모 노드로 바꿔준다.
        parent[y] = x
        # y를 탐색한 횟수도 x에 더해준다.
        visited[x] += visited[y]


# 부모 노드 찾기
def find(target):
    # 자신이 부모 노드이면 자기 자신을 리턴
    if target == parent[target]:
        return target

    # 재귀 함수를 통해 부모 노드 찾기
    else:
        parent[target] = find(parent[target])
        return parent[target]


t = int(sys.stdin.readline())
for _ in range(t):
    f = int(sys.stdin.readline())
    parent = dict() # 딕셔너리형
    visited = dict() # 딕셔너리형

    # 친구 관계의 수만큼 반복하여 친구의 관계를 확인
    for i in range(f):
        a, b = map(str, sys.stdin.readline().split())

        # 친구 관계의 없는 친구이면 추가해준다.
        # 탐색 횟수도 카운트한다.
        if a not in parent:
            parent[a] = a
            visited[a] = 1

        if b not in parent:
            parent[b] = b
            visited[b] = 1

        # 두 친구의 관계을 합친다.
        union(a, b)

        print(visited[find(a)])