백준 4195번 : 친구 네트워크 - 파이썬, Python 풀이

들어가며

오랜만에 백준 문제를 올린다. 이번 문제는 유니온 파인드(Union Find), 디스조인트 셋(Disjoint Set)이라고 불리는 알고리즘의 대표적인 문제인 친구 네트워크라는 문제이다.

현재 solved.ac 기준 골드 2에 랭크돼있다. 골드 2인 이유는 유니온 파인드를 응용할 아이디어를 생각해내야 하고 문자열 처리를 해야 하기 때문이다. 아이디어만 생각해낸다면 파이썬으론 문자열 처리가 무척 쉬우므로 알고리즘 적용 시 무리 없이 풀린다.

풀이

우리는 입력된 두 사람의 친구 네트워크를 구성했을 때(유니온 파인드를 통해 트리를 구성했을 때), 친구의 수(트리를 구성하는 노드의 수)를 구해주어야 한다. 그러므로 단순히 노드끼리 합쳐 줄 것이 아니라, 동시에 자손 노드 자기 자신과 거느리고 있는 노드 수를 부모 노드에 전달해주어야 한다.

노드마다 부모 노드를 저장할 parent 딕셔너리와 노드마다 자기 자신을 포함하여 거느리고 있는 노드 수를 저장할 cnt 딕셔너리를 선언한다. 입력이 들어올 때마다 유니온파인드를 진행하여 친구 네트워크를 구성하고 네트워크 내 노드 수를 출력해야 한다.

그렇기 때문에 입력마다 parent 딕셔너리와 cnt 딕셔너리에 입력된 친구에 대한 정보의 존재여부를 확인하고 없다면 새롭게 등록해주고 유니온 파인드를 진행해야한다.

cnt = dict()
parent = dict()
F = int(sys.stdin.readline())
for _ in range(F):
    temp = sys.stdin.readline().split()
    if temp[0] not in cnt: # cnt에 temp[0]에 있는 문자열에 대한 정보가 있는지 확인
        cnt[temp[0]] = 1
        parent[temp[0]] = temp[0]
    if temp[1] not in cnt: # cnt에 temp[1]에 있는 문자열에 대한 정보가 있는지 확인
        cnt[temp[1]] = 1
        parent[temp[1]] = temp[1]

merge함수 구성시 기본적인 merge 함수 구성 외에도 부모노드에 자손노드 자기자신을 포함하여 거느린 노드들의 수를 더해주는 코드를 넣어주자.

def merge(node1, node2):
    node1 = find(node1)
    node2 = find(node2)

    if node1 == node2:
        return node1

    parent[node1] = node2
    cnt[node2] += cnt[node1] # 부모노드에 노드 수를 누적

또한, 출력 시 노드가 루트 노드를 갱신할 수 있도록 출력 이전 find() 함수를 적용해주자.

print(cnt[find(temp[0])])

전체 코드

import sys

def find(node):
    if node == parent[node]:
        return node
    parent[node] = find(parent[node])
    return parent[node]

def merge(node1, node2):
    node1 = find(node1)
    node2 = find(node2)

    if node1 == node2:
        return node1

    parent[node1] = node2
    cnt[node2] += cnt[node1]

t = int(sys.stdin.readline())

for _ in range(t):
    cnt = dict()
    parent = dict()
    F = int(sys.stdin.readline())
    for _ in range(F):
        temp = sys.stdin.readline().split()
        if temp[0] not in cnt:
            cnt[temp[0]] = 1
            parent[temp[0]] = temp[0]
        if temp[1] not in cnt:
            cnt[temp[1]] = 1
            parent[temp[1]] = temp[1]

        merge(temp[0], temp[1])
        print(cnt[find(temp[0])])