[알고리즘/Union-Find/Python] 백준 4195 - 친구 네트워크

문제 : 4195번: 친구 네트워크 (acmicpc.net)

4195번: 친구 네트워크

사용언어 : Python3

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

 

예제 입력

2
3
Fred Barney
Barney Betty
Betty Wilma
3
Fred Barney
Betty Wilma
Barney Betty

 

출력

 친구 관계가 생길 때마다, 두 사람의 친구 네트워크에 몇 명이 있는지 구하는 프로그램을 작성하시오.

예제 출력

2
3
4
2
2
4

 

사용된 알고리즘 : Union - Find 알고리즘

 상호 배타적 집합 (Disjoint Set) 을 결합하거나, 어떤 두 노드가 같은 집합에 있는지 확인이 필요할 때 사용되는 알고리즘입니다. 이를 위해 두 노드의 집합을 결합하는 Union연산과 특정 노드가 포함된 집합을 찾는 Find 연산 두가지의 연산이 존재합니다.

 

• Find
  • 노드 x 가 포함된 집합을 찾는 연산
• Union
  • 노드 x와 노드 y가 포함된 집합을 합치는 연산

 

예시 )

1. Fred Barney

parent
Fred	Fred
Barney	Barney

number
Fred	1
Barney	1

↓

parent
Fred	Fred
Barney	Fred

number
Fred	2
Barney	2

 

 

코드

 파이썬에서 집합을 구현하기 위해 딕셔너리를 사용하였다. parent 딕셔너리와 number 딕셔너리를 생성한다. 

 

 Fred Barney를 입력 받았을 경우 Fred와 Barney를 각각 x, y로 받고 키로하여 parent와 number 딕셔너리에 각각 초기화해준다. 그 후 x, y를 union 연산한다.

if x not in parent:
  parent[x] = x
  number[x] = 1
if y not in parent:
  parent[y] = y
  number[y] = 1
union(x, y)

 

 유니온 연산은 x와 y를 find 연산한 후 x와 y로 결과값을 반환받는다. 만약 이 값이 같지 않다면 parent[y]의 값을 x로 교체하고, number[x]에는 number[y]의 값을 더해준다.

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

  if x != y:
    parent[y] = x
    number[x] += number[y]

 

 파인드 연산은 재귀함수로서 parent 값이 자기자신인 경우와 그렇지 않은 경우로 나뉜다.

 parent 값이 자기 자신인 경우에는 자신을 리턴한다. 그렇지 않은 경우엔 재귀호출로 자신의 부모인 p를 찾고 parent의 자신의 값을 p로 지정한다. 그후 parent 값을 리턴한다.

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

 

파이썬 코드

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

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

  if x != y:
    parent[y] = x
    number[x] += number[y]

tC = int(input())

for _ in range(tC):
  parent = dict()
  number = dict()

  F = int(input())
  for _ in range(F):
    x, y = input().split()

    if x not in parent:
      parent[x] = x
      number[x] = 1
    if y not in parent:
      parent[y] = y
      number[y] = 1
    union(x, y)  
    print(number[find(x)])