[Python Algorithm] 백준 친구 네트워크 - 4195

풀이 과정

유니온 파인드 알고리즘의 응용문제이다.
가장 먼저 든 생각은 전체 네트워크에 몇명이 존재하는지 어떻게 구해야 하는지에 대해서였다.

그런데 잘 생각해보니 유니온 파인드 알고리즘에서 find 함수는 항상 노드의 가장 위쪽 부모노드를 찾아주기 때문에 그냥 하나의 dict를 더 선언해서 그곳에 값을 잘 저장만 해주면 해결이 될 것 같았다.

그래서 union을 해줘야하는 경우(find함수 이후 부모가 다른경우)에 항상 첫번째 이름(name1)이 두번째 이름(name2)의 부모가 되게 지정해주고, 네트워크의 개수를 저장하는 dict는 항상 name1이 업데이트 되도록 구현해주었다.

해당 부분의 코드는 다음과 같다.

  if name1 != name2:
    parents[name2] = name1
    cnt[name1] = cnt[name2] + cnt[name1]

풀이에 대한 결정적인 부분은 해결하였지만 어째서인지 바로 해결하진 못했는데, 그 이유는 두개의 이름을 입력받고 dict에 넣어주는 부분이었다.
처음에는 아래 코드와 같이 입력을 받으면 parents[name] = name 으로 바로 초기화 해주었는데 이렇게 되면 중복된 이름이 하나만 들어와도 이전 트리가 망가지게 되서 문제가 풀리지 않는다.

 name1, name2 = map(str, input().split(' '))
 	parents[name1] = name1
    parents[name2] = name2

해당 이슈까지 해결한 이후에야 겨우 통과를 할 수 있었다.

---

최종 코드

from collections import defaultdict

T = int(input())

def find(parents, name):
  if name != parents[name]:
    parents[name] = find(parents, parents[name])
  return parents[name]

def union(parents, cnt, name1, name2):
  name1 = find(parents, name1)
  name2 = find(parents, name2)

  if name1 != name2:
    parents[name2] = name1
    cnt[name1] = cnt[name2] + cnt[name1]
  print(cnt[name1])

for _ in range(T):
  cnt = defaultdict(lambda : 1)
  parents = {}

  N = int(input())

  for i in range(N):
    name1, name2 = map(str, input().split(' '))
    if not name1 in parents:
      parents[name1] = name1
    if not name2 in parents:
      parents[name2] = name2

    union(parents, cnt, name1, name2)
    

---

결과

---

Link

친구 네트워크 문제 출처
GitHub 코드