백준 1717 집합의 표현 파이썬 문제풀이 (python 유니온파인드)

문제 보러가기

 

백준 집합의 표현 파이썬 문제풀이

유니온 파인드 문제에 대해서는 많이 공부하지 않아서

개념을 잘 봐두려고 한다. 

 

유니온파인드는 요소가 어느 집합에 포함되는지를 알아내는 방법인데. 

요소의 부모 리스트를 담고, 최종적으로 루트 노드가 일치하면 같은 집합으로 판별하는 방법이다. 

 

문제에서는 0, a ,b 는 합집합

1, a, b 는 같은 집합인지 확인한다.

 

 

N,M = map(int, input().split())

par = [i for i in range(N+1)]
# [0,1,2,3,4,5,6,7]

def _union(A,B):
    par[B] = A

def _find(A):
    if par[A] == A: # 만약에 스스로를 부모로 칭하고 있다면!
        return A
    else:
        return _find(par[A])

for _ in range(M):
    X,A,B = map(int, input().split())

    if X == 0:
        _union(A,B)
    else:
        if _find(A) == _find(B):
            print("YES")
        else:
            print("NO")

0으로 입력했을 때 union 함수를 통해서 같은 부모를 공유하게 하여

같은 부모를 공유하는 집합으로 만들어주고.

 

1을 입력한 경우 두 요소가 같은 집합인지 find 함수를 실행한다.

find 함수는 계속 부모를 향해 재귀를 돌기 때문에 결과적으로 루트 노드를 비교하고.

같은 루트 노드라면 같은 집합임을 확인하는 로직이다.

 

하지만 결과는 틀렸다..

어디서 문제일까. 

 

최적화

import sys
sys.setrecursionlimit(1000000)
N,M = map(int, input().split())

par = [i for i in range(N+1)]
# [0,1,2,3,4,5,6,7]
rank = [0 for _ in range(N+1)]

def _union(A,B):
    A = _find(A)
    B = _find(B)

    if A == B:
        return

    if rank[A] < rank[B]:
        par[A] = B
    elif rank[B] < rank[A]:
        par[B] = A
    else:
        par[A] = B
        rank[B] +=1

    par[B] = A

def _find(A):
    if par[A] == A: # 만약에 스스로를 부모로 칭하고 있다면!
        return A
    else:
        par[A] = _find(par[A])
        return _find(par[A])

for _ in range(M):
    X,A,B = map(int, input().split())

    if X == 0:
        _union(A,B)
    else:
        if _find(A) == _find(B):
            print("YES")
        else:
            print("NO")

어차피 부모 노드를 타고 올라가서 루트 노드는 같다면, 

부모 노드에 find한 루트를 넣어서 재귀경로를 절약하는 최적화를 할 수 있다.

 

또한 union 함수의 경우도 랭크를 이용해서 부모의 랭크를 증가시키고

랭크가 높은쪽을 부모로 삼도록 하여 최적화를 진행했다.

 

하지만 메모리 초과로 실패했는데.

뭐가 또 문제가 되는 걸까.

 

메모리를 초과했다는 것은 너무 배열을 낭비했다는 뜻.

rank 리스트를 사용하지 않고 같은 효과를 내도록 해보자.

 

import sys
sys.setrecursionlimit(10000)
N,M = map(int, input().split())

par = [i for i in range(N+1)]
# [0,1,2,3,4,5,6,7]
rank = [0 for _ in range(N+1)]

def _union(A,B):
    A = _find(A)
    B = _find(B)

    if A == B:
        return

    if A > B:
        par[A] = B
    else:
        par[B] = A

    par[B] = A

def _find(A):
    if par[A] == A: # 만약에 스스로를 부모로 칭하고 있다면!
        return A
    else:
        par[A] = _find(par[A])
        return _find(par[A])

for _ in range(M):
    X,A,B = map(int, input().split())

    if X == 0:
        _union(A,B)
    else:
        if _find(A) == _find(B):
            print("YES")
        else:
            print("NO")

rank를 사용하지않고 무조건 숫자가 작은 쪽이 부모가 되도록 했다. 

 

그럼에도 메모리 초과..

 

정답 코드 

import sys
sys.setrecursionlimit(1000000)
input = sys.stdin.readline
N,M = map(int, input().split())

par = [i for i in range(N+1)]

def _union(A,B):
    A = _find(A)
    B = _find(B)

    if A == B:
        return

    if A < B:
        par[A] = B
    else:
        par[B] = A

def _find(A):
    if par[A] == A: # 만약에 스스로를 부모로 칭하고 있다면!
        return A
    else:
        par[A] = _find(par[A])
        return _find(par[A])

for _ in range(M):
    X,A,B = map(int, input().split())

    if X == 0:
        _union(A,B)
    else:
        if _find(A) == _find(B):
            print("YES")
        else:
            print("NO")

문제는 union의 최적화 로직에 있었다.

마지막에 다시 A를 B의 부모로 넣는 로직 때문에 최적화가 무효가 되었었는데.

 

수정하고 나니. 정답이다.

 

그리고 참고로 백준에서 pypy3로 하면 왜인지 정답처리가 안되어서

python3로 해야 정답 처리가 되었다.

 

또 파이썬에서 연산을 줄이는 설정인

input = sys.stdin.readline

이 설정을 해주지 않으면 정답이 틀리진 않더라도 채점속도가 무진장 느리다.

(아래쪽이 input설정 없이 채점한 결과)