4195. 친구 네트워크 (Python)

4195번: 친구 네트워크

 

풀이

 유니온 파인드를 사용하여 해결할 수 있다. 해당 문제에서는 2개의 딕셔너리를 사용한다. 하나는 친구 관계에서 루트가 될 사용자를 설정하는 parent 딕셔너리이고, 나머지 하나는 친구가 몇 명이 있는지를 카운트하는 friendCount 딕셔너리이다. 

 

 친구 관계가 주어졌을 때 처음 등장하는 사용자라면 friendCount에 사용자 아이디를 키 값으로 하여 1로 설정해준다. 그 후 두 사용자를 union 함수를 사용하여 관계를 설정해준다. 만약 두 사용자의 루트가 같은 사람이라면 하나의 친구 네트워크에 있기 때문에 따로 설정해주지 않아도 되지만 다르다면 두 개의 친구 네트워크를 하나로 묶어줘야 한다. 관계 설정 후 firendCount 딕셔너리를 사용하여 총 몇 명의 사용자가 있는지를 출력해주면 된다.

 

코드

from collections import defaultdict
import sys

input = sys.stdin.readline


# name 사용자에 대한 루트 사용자를 찾기 위한 함수
def find(name):
    if parent[name] == '':
        return name
    parent[name] = find(parent[name])
    return parent[name]


# n1과 n2의 관계 설정을 위한 함수
def union(n1, n2):
    n1 = find(n1)  # n1 사용자의 루트 사용자를 find
    n2 = find(n2)  # n2 사용자의 루트 사용자를 find

    if n1 == n2:  # 두 사용자의 루트 사용자가 같다면 이미 하나의 네트워크에 속한 것이므로 따로 설정 필요 x
        return

    parent[n2] = n1  # n2 사용자의 루트 사용자를 n1의 루트 사용자로 설정
    friendCount[n1] += friendCount[n2]  # 친구 숫자 업데이트


if __name__ == '__main__':
    T = int(input())  # 테스트 케이스의 개수

    for _ in range(T):
        F = int(input())  # 친구 관계의 수
        parent = defaultdict(str)  # parent['name']: 'name' 사용자에 대한 루트 사용자의 아이디
        # friendCount['name']: 'name' 사용자의 친구 네트워크에 몇 명이 있는가를 의미
        friendCount = defaultdict(int)

        for i in range(F):
            n1, n2 = map(str, input().split())

            if friendCount[n1] == 0:  # n1이 처음 등장할 경우 친구의 수를 1로 설정
                friendCount[n1] = 1
            if friendCount[n2] == 0:  # n2가 처음 등장할 경우 친구의 수를 1로 설정
                friendCount[n2] = 1

            union(n1, n2)  # n1과 n2 사용자의 관계 설정
            name = find(n1)  # 둘 중 한 사용자의 루트 사용자를 찾고
            print(friendCount[name])  # 그 루트 사용자의 친구의 수를 출력한다