[파이썬] 백준 1717 집합의 표현

BFS로 구현을 했다. 테스트 케이스는 잘 나오고 당연히 맞은 줄 알았는데 시간 초과가 나왔다. 구글링 해보니 union-find 라고 한다. 

from collections import deque
from itertools import combinations
import sys
 
def check(k,a,b):
 
    # a 기준으로 찾고 거기에 b가 있으면 true 반환
 
    queue = deque()
 
    visited = []
 
    for i in range(0,len(k)):
        if a in k[i]:
            queue.append(k[i])
 
    while queue:
 
        if len(queue)==0:
            break
 
        x,y = queue.popleft()
 
        if x not in visited:
            visited.append(x)
 
        if y not in visited:
            visited.append(y)
 
 
        for j in range(0,len(k)):
            if (x in k[j] or y in k[j]):
 
                if x not in visited and y not in visited:
                    queue.append([x,y])
 
 
    if b in visited:
        return True
    else:
        return False
 
 
 
n, m = map(int,input().split(" "))
 
k = [ ]
 
for i in range(0,m):
 
    checks, a, b  = map(int,input().split(" "))
 
    # check == 0
 
    if checks == 0:
        k.append([a,b])
 
 
 
 
 
    # checks == 1
 
    if checks == 1:
        if check(k,a,b) == True:
            print("YES")
        else:
            print("NO")
 
 
 
 
 
 
 

 

아래가 정답 코드라고 한다. 얼추 이해는 되는데 자세한 건 다시 봐야겠다.

import sys
input = sys.stdin.readline
sys.setrecursionlimit(100000)
n, m = map(int, input().split())
 
def union(x, y):
    x = find(x)
    y = find(y)
    if x == y:
        return
    parent[x] = y
    return
 
def find(x):
    if parent[x] == x:
        return x
    parent[x]=find(parent[x])
    return parent[x]
 
parent = [0]*(n+1)
for i in range(n+1):
    parent[i] = i
 
for i in range(m):
    x, a, b = map(int, input().split())
    if x == 0:
        union(a,b)
    else:
        pa = find(a)
        pb = find(b)
        if pa == pb:
            print('YES')
        else:
            print('NO')