[BOJ] 4195 친구 네트워크

https://www.acmicpc.net/problem/4195

1차 시도도 안 함

처음에는 인접리스트나 뭐 그런거 만들어서 DFS나 BFS같은 걸로 돌려야 하나 싶었다. 그런데 마음속에서 생긴 DFS코드를 만들어야하는 거부감이 생겨서 고민하다가 DFS로도 안 될 것 같아서, 시간제한을 50분을 줬는데 그래도 안 돼서 포기하고 해설 강의를 들었다.

Union Find

해설강의에서는 Union Find를 통해서 문제를 풀면 된다고 했다. 먼저 단순한 Union Find 코드를 보여줬는데 다음과 같다.

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)

  parent[y] = x

parent = []
for i in range(0, 5):
  parent.append(i)

union(1, 4)
union(2, 4)

for i in range(1, len(parent)):
  print(find(i), end = ' ')
  

먼저 부모 자식 관계를 표현할 배열을 하나 만들고 그러고 자기 자신을 부모로 초기활르 시킨다. 그러고 나서 엣지가 하나 들어오면 왼쪽에 있는 애를 부모로 삼게 하는데 그 과정은 일단 각자의 부모를 찾은 뒤에 왼쪽에 있는 값을 오른쪽에 있는 값이 부모로 삼게 하는 방식이다.
부모 찾는 과정은 자기 자신이 부모면 바로 반환하고, 아니라면 부모의 부모를 재귀함수를 통해서 찾게 한다. 이 부분이 이해가 잘 안 갔는데 샘플위 수가 적어서 내가 착각한 것이었다. 부모의 부모까지 그러고 나중에 그 위에 아무도 없는 값이 나오게 될 때까지 재귀함수는 들어가고 나중에 같은 연산을 안하게 하려고 마지막 조상을 찾게 되면 걔를 자신의 부모로 삼게 하고 제일 마지막에는 최종 조상을 반환해서 다른 네트워크와 비교하게 하는 방식이다.

코드

import sys
input = sys.stdin.readline

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()

  for _ in range(int(input())):
    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)])
    

위의 Union Find 코드에다가 네트워크 안에 몇 명있는지 확인하는 사전을 하나 추가하는 코드만 추가해서 해결했다. 처음에 시간초과 나길래 보니까 입력값이 최대 10만이어서 sys의 readline을 가져와서 해결해줬다.

예전에 다익스트라 공부하다가 Union Find를 발견했을 때는 귀찮아보였는데 강의 들으면서 하나하나 따라가니까 생각보다 단순하고 코드 양도 얼마 안 돼서 여러군데 써먹을 수 있을 것 같다.