[백준 4195] 친구 네트워크 (Python) / G2

문제

4195. 친구 네트워크

코드

t = int(input())

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


def union(a,b):
    a = find(a)
    b = find(b)

    # a와 b의 부모가 같지 않으면 a에 b 연결!
    if a != b:
        parent[b] = a
        number[a] += number[b]

    print(number[a])



for _ in range(t):
    f = int(input())
    # 해당 key의 부모를 value로 넣어주는 dict
    parent = {}
    # 해당 키와 연결되어 있는 친구 네트워크 수를 저장해주는 dict
    number = {}

    for i in range(f):
        a, b = input().split()
        # 만약 parent에 a 키값이 없으면
        if a not in parent:
            parent[a] = a
            number[a] = 1
        if b not in parent:
            parent[b] = b
            number[b] = 1

        # 친구 연결이므로 union-find 알고리즘 사용
        union(a,b)

접근 방법

친구 연결을 하는 문제이므로 Union-Find 알고리즘을 사용해야된다고 판단하였다.

기본 Union-Find 알고리즘과 다른 것이, 연결 이름이 String 형태이므로 리스트가 아닌 딕셔너리 형태롤 값을 받아 부모 노드를 value로 연결하는 식으로 구현하였다.

number이라는 딕셔너리에는 key에 해당하는 사람이 연결된 친구 수를 value로 넣어주었다.

기본 Union-Find처럼 트리의 최상위 부모는 자기자신을 value값으로 갖게해서 union 함수와 find 함수를 작성해주었다.

그리고 union함수에서 두 값의 부모가 같지 않으면 서로 연결시켜주는 부분에서 최적화시키지 않고 그저 a에 b를 연결시켜주는 방식으로 구현하여 주었다.

추가

제출에 시간이 오래걸려서

import sys
input = sys.stdin.readline

을 추가하여 주었더니 시간차이가 생각보다 많이 났다.

input을 stdin을 사용하여 시간복잡도가 얼마나 줄어드는지 공부가 필요해보임!