문제

민혁이는 소셜 네트워크 사이트에서 친구를 만드는 것을 좋아하는 친구이다. 우표를 모으는 취미가 있듯이, 민혁이는 소셜 네트워크 사이트에서 친구를 모으는 것이 취미이다.

어떤 사이트의 친구 관계가 생긴 순서대로 주어졌을 때, 두 사람의 친구 네트워크에 몇 명이 있는지 구하는 프로그램을 작성하시오.

친구 네트워크란 친구 관계만으로 이동할 수 있는 사이를 말한다.

입력

첫째 줄에 테스트 케이스의 개수가 주어진다. 각 테스트 케이스의 첫째 줄에는 친구 관계의 수 F가 주어지며, 이 값은 100,000을 넘지 않는다. 다음 F개의 줄에는 친구 관계가 생긴 순서대로 주어진다. 친구 관계는 두 사용자의 아이디로 이루어져 있으며, 알파벳 대문자 또는 소문자로만 이루어진 길이 20 이하의 문자열이다.

코드

import sys
input = sys.stdin.readline
tc = int(input())

def find(x):
    if x == parent[x]:
        return x
    else:
        root_x = find(parent[x])
        parent[x] = root_x
        return parent[x]

def union(x, y) :
    root_x = find(x)
    root_y = find(y)

    if root_x != root_y:
        parent[root_y] = root_x
        number[root_x] += number[root_y]


for i in range(tc):
    n = int(input())
    parent = dict()
    number = dict()

    for j in range(n):
        a, b = map(str, input().split())

        if a not in parent:
            parent[a] = a
            number[a] = 1

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

        union(a,b)
        print(number[find(a)])

설명

처음엔 딕셔너리 하나를 생성해서 각 이름에 어떤 친구가 있는지 저장하고 인접리스트 처럼 하나하나 visited = True처럼 방문 처리를 해주었다. 하지만 시간초과가 계속 나서 결국 찾아보았는데 union-find를 사용해서 푸는 문제였다. 입력 받은 이름들이 처음 들어오는 이름이라면 parent에 첫 생성을 해주고 두 이름을 서로 union해준다. 그리고 출력할때는 find함수를 사용해서 최상단 부모를 가르키는 이름들의 숫자를 출력해주었다.