집합의 표현

집합 ‘set()’을 직접 사용해서 풀면 시간초과 혹은 메모리초과가 발생하는 문제이다.

각 숫자가 가르키는 다음 숫자(포인터)를 활용해서 같은 집합에 속하는지 혹은 같은 집합에 속하지 않는지 알 수 있다. 포인터가 향하는 곳을 따라 탐색하다가, 포인터가 자기 자신을 가리키는 숫자가 나타나는 경우 그 숫자가 root-node이다.

첫 번째 표 1,3,4가 같은 root-node를 가지고 있다.(같은 집합에 속한다.)

• 0 ~ 6의 숫자들이 다음과 같은 값들을 가지고 있다.
• 1은 자신을 가리키고 있음으로 root-node는 1이다.
• 3은 1을 가리키고 있고, 1은 1을 가리키고 있다. 따라서 3의 root-node는 1이다.
• 4또한 1을 가리키고 있고, 1은 자기자신을 가리키고 있다. 따라서 4의 root-node는 1이다.

두 번째 표 - 1과 5가 속한 집합을 합쳐서 1이 가리키는 포인터가 5로 변경되었다. 1, 3 그리고 4의 root-node는 어떻게 바뀌었는지 알아보자.

• 1은 5를 가리키고, 5는 자기자신을 가리키고있다. 1의 root-node가 5로 바뀌었다.
• 3은 1을 가리키고, 1은 5, 5는 5를 가리키고 있다. 따라서 3의 root-node도 5로 바뀌었다.
• 4또한 3과 동일하다.

이렇게 1,3,4,5가 같은 집합에 속하게 되었음을 알 수 있다.

위의 과정을 코드로 구현하기 위해, root-node를 찾는 find-root함수와 각 집합을 합하는 union합수를 만들었다.

#재귀 함수를 이용해 root 노드를 찾아간다.
def find_root(x):
    if x == root[x]:
        return x
    else:
        #밑의 union 함수에서 같은 root Node를 가지는 모든 원소들의 값을 업데이트하는 것이 아니기 때문에
        #root Node를 찾아가는 과정에서 이를 다시 업데이트 해준다.
        root[x] = find_root(root[x])
        return root[x]

def union(x, y):
    a = find_root(x)
    b = find_root(y)
    #root Node의 값을 변경해준다. 이때, 하위 원소들의 root Node값들은 별도로 업데이트 하지 않는다.
    if a != b:
        root[b] = a

TestCase

input:
6 6
0 1 3
0 3 4
0 4 5
1 5 1
0 6 5
1 6 4

answer:
yes
yse

위의 TestCase에 대해 각 숫자들의 포인터를 어떻게 바꿔나가는지 그림으로 나타내 보았다.

0 1 3

• 1과 3을 합쳐야 한다.
• 1의 root-node는 자기자신인 1이다.
• 3의 root-node는 자기자신인 3임으로, 3이 1을 가리키도록 값을 변경한다.
• 3의 root-node가 1로 변경되어 1과 3이 같은 집합에 속하게 되었다.

0 3 4

• 3의 root-node는 1이다.
• 4의 root-node는 4이다. 4가 3의 root-node인 1을 가리키도록 값을 변경한다.
• 4의 root-node가 1로 변경디어 3과 4도 같은 집합에 속하게 되었다.
• 현재 1, 3, 4가 같은 root-node를 가지고 있다.

0 5 1

• 5의 root-node는 5이다.
• 1의 root-node는 1이다. 1이 5의 root-node를 가리키도록 값을 5로 업데이트해준다.
• 1의 root-node가 5로 변경되었다. 이때, 3과 4의 root-node는 어떻게 바뀌었을지 살펴보자.
  • 3은 1을 가리키고, 1은 5를 가리키고있다. 5는 자기자신을 가리키고 있음으로, 3의 root-node도 5로 변경되었음을 알 수 있다.
  • 4 또한 마찬가지이다.

1 5 3

• 5의 root-node는 5이다.
• 3의 root-node는 5이다. 이때, 탐색하는 과정에서 3의 root-node가 5임을 알게 되었음으로, 값을 1에서 5로 업데이트 해준다.(필수적인 과정은 아니지만, 이렇게 업데이트 함으로써 시간복잡도를 줄일 수 있다.)
• root-node가 동일함으로 같은 집합에 속해있음을 알 수 있다. ‘YES’출력

0 6 5

• 6의 root-node는 6이다.
• 5의 root-node는 5이다. 5가 6의 root-node를 가리키도록 값을 6으로 변경한다.

1 6 4

• 6의 root-node는 6이다.
• 4는 1을 가리키고, 1은 5를 가리키고, 5는 6을 가리킨다. 6은 자기자신을 가리킴으로 4의 root-node는 6이다.
• 이때, 1과 4의 root-node가 6임을 알았음으로, 1과 4의 root-node를 업데이트해준다.
• root-node가 동일함으로 같은 집합에 속해있음을 알 수 있다. ‘YES’출력

CODE

import sys
read = sys.stdin.readline

n, m = map(int, read().split())
root = list(i for i in range(n+1))

#재귀 함수를 이용해 root 노드를 찾아간다.
def find_root(x):
    if x == root[x]:
        return x
    else:
        #밑의 union 함수에서 같은 root Node를 가지는 모든 원소들의 값을 업데이트하는 것이 아니기 때문에
        #root Node를 찾아가는 과정에서 이를 다시 업데이트 해준다.
        root[x] = find_root(root[x])
        return root[x]


def union(x, y):
    a = find_root(x)
    b = find_root(y)
    #root Node의 값을 변경해준다. 이때, 하위 원소들의 root Node값들은 별도로 업데이트 하지 않는다.
    if a != b:
        root[b] = a


for _ in range(m):
    op, a, b = map(int, read().split())
    if op == 0:
        union(a, b)
    else:
        if find_root(a) == find_root(b):
            print('YES')
        else:
            print("NO")

틀린 부분이 있을 수 있습니다. 피드백 주시면 고치도록 하겠습니다! 감사합니다.👍

[꼭 다시 풀어보기]