[백준] Union & Find - 친구 네트워크

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

 

Union-Find는 크게 두가지 테스크를 진행 할 수 있다고 언급을 했다.

1. Cycle 여부 확인

2. 동일 그래프 소속 여부

 

크게 두가지를 알 수 있는데 지난번 테스크는 사이클 여부를 체크하는 문제를 풀어봤다.

이번에는 동일 그래프에 속하는지 판단 + @로 조금 다른 find를 하는 문제이다.

 

이번 알고리즘은 두 사람의 이름이 "A B"와 같이 주어지면, A와 B는 친구 네트워크가 생성된다는 의미이다.

이렇게 되면 단순하게 노드를 연결하는 방법(Union)이 된다. 즉, Union 방식으로 그래프를 생성해 나가면 된다.

 

하지만 한가지 다른점은 Find가 친구관계가 input으로 들어왔을때, 해당 그래프는 몇명의 친구와 연관이 있는지 바로바로 출력을 해야한다.

 

이렇게 되면, 기존 Union-Find를 사용하여 계속해서 몇명이 연결되어있는지 Find를 하게되면 속도의 문제로 이 문제는 풀 수 없는 상황이 된다.

 

이걸 해결하기 위해서 dict를 이용하여 Union할때 각 그래프별 연결된 사람의 수를 지속적으로 업데이트하여 출력을 최대한 빠르게 하기 위해 작성을 하였다.

 

이 문제를 풀기위해 2틀정도의 시간이 걸렸는데, 시간이 오래 걸린 이유가 따로있었다.

Union-Find는 딕셔너리 형태로 잘 작성하여 자체 테스트해봤을때는 아무 문제가 없었다.

 

이럴땐 아주 사소한 것이 문제를 일으키는 경우가 상당히 많다.

필자는 여기서 input을 sys.stdin.readline을 사용해여 입력 속도를 빠르게 진행하였다.

 

이게 여기서는 문제를 발생시켰다.

 

내용을 찾아보니 위와같은 내용을 찾을 수 있었다.

즉, input()으로 진행을 하면 '줄 바꿈 제거'의 과정이 들어가있어서 첫날 푼 코드로도 통과가 되었다...

반면 sys.stdin.readline은 그런게 없고, 입력한 모든 문자열을 그대로 받아들이기 때문에 .strip()을 통해 빈칸을 제거해줘야 옳바르게 사용할 수 있다는 점을 알았다..

 

다음에는 반드시 .strip()을 사용하도록 해야겠다..

 

 

def get_parent(store_dict , n ):
  '''
    node의 root parent를 찾는다.
  '''
  # parnet와 현재 node가 같으면 root
  # print(store_dict)
  # print(n)
  if store_dict[n] == n:
    return n
  store_dict[n] = get_parent(store_dict , store_dict[n])
  return store_dict[n]

def union_(store_dict , count_dict , n1, n2 ):
  pn1 = get_parent(store_dict , n1)
  pn2 = get_parent(store_dict , n2)

  if pn1 < pn2 :
    store_dict[pn2] = pn1
    count_dict[pn1] += count_dict[pn2]
    
    print(count_dict[pn1])
     
  elif pn1 > pn2 :
    store_dict[pn1] = pn2
    count_dict[pn2] += count_dict[pn1]
    
    print(count_dict[pn2])
  else:
    print(count_dict[pn1])
    
import sys

sys.setrecursionlimit(10**6)
input = sys.stdin.readline


T = int(input().strip())
for _ in range(T):
  number = int(input().strip())

  store_dict = {}
  count_dict = {}

  for _ in range(number):
    A, B = input().strip().split(" ")


    index = 0
    if A not in store_dict:
      store_dict[A] = A 
      count_dict[A] = 1
      index += 1 

    if B not in store_dict:
      store_dict[B] = B 
      count_dict[B] = 1
      index += 1 

    
    union_(store_dict , count_dict , A, B )