BAEKJOON #4195) 친구 네트워크 | union-find | Python

4195번: 친구 네트워크

union find 알고리즘을 이용해서 푸는 문제입니다. union-find에 관한 자세한 내용은 아래 글을 참고해주세요:)

Union Find 알고리즘 | 개념과 Python 코드 구현

 

처음에는 find 함수를 아래처럼 작성하여 제출하였더니 RecursionError가 났습니다. 재귀를 너무 많이 돌아서 생기는 에러입니다. 그래서 find함수에서 경로 압축을 해주어야 시간복잡도를 줄일 수 있습니다. (위의 Union Find 알고리즘 글에서 더 자세한 설명을 보실 수 있습니다.)

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

 

# -*- coding: utf-8 -*-
# boj 4195 친구 네트워크

import sys
input = sys.stdin.readline

test_case = int(input())

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


def union(x, y):
	x = find(x)
	y = find(y)
	if x != y:
		parent[y] = x
		network_size[x] += network_size[y]


for _ in range(test_case):
	parent = dict()
	network_size = dict()

	f = int(input())
	for __ in range(f):
		x, y = input().strip().split(' ')

		if x not in parent:
			parent[x] = x
			network_size[x] = 1
		if y not in parent:
			parent[y] = y
			network_size[y] = 1

		union(x, y)
		print(network_size[find(x)])

최종적으로 제출하여 정답으로 처리된 코드는 다음과 같습니다.

parent와 network_size라는 딕셔너리 2개를 만들어서 부모를 저장하는 사전과 네트워크 사이즈를 저장하는 사전으로 사용합니다.