[백준 1717 🥇] 집합의 표현

 

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

1717번: 집합의 표현

문제풀이

import sys
sys.setrecursionlimit(100000)

n,m=map(int, sys.stdin.readline().split())

# union-find 알고리즘을 사용
# 집합의 root 번호를 나타내는 배열 사용
root_set=[x for x in range(n+1)]

def find(num):
    global root_set
    if num!=root_set[num]:
        root_set[num] = find(root_set[num])
    return root_set[num]
    # if num==root_set[num]:
    #     return num
    # else:
    #     root_set[num]=find(root_set[num])

def union(x,y):
    global root_set
    x=find(x)
    y=find(y)
    root_set[y]=x


for _ in range(m):
    command, n1, n2=map(int, sys.stdin.readline().split())
    if command==0:
        # 두 번호가 속한 집합을 합친다.
        union(n1,n2)
    elif command==1:
        # 두 번호가 같은 집합에 속하는지 확인한다.
        if find(n1)==find(n2):
            print('YES')
        else:
            print('NO')

집합을 합치고, 집합의 루트를 찾는 Union-Find 알고리즘을 사용하는 문제였다. 

집합을 트리 형태로 구현해서 집합끼리 합쳐줄 땐, 같은 루트를 가지는 트리로 만들어주고,

동일한 집합에 속하는지 찾을 때는 동일한 루트를 가지는지 확인해주면 된다. 

모든 집합의 루트를 자기 자신으로 초기화해준다.

index	0	1	2	3	4	5	6	7
	0	1	2	3	4	5	6	7

그리고 어떤 숫자가 속한 집합의 루트를 찾아주는 find 함수는 루트를 찾는 숫자의 root_set[] 배열 값이 자기 자신과 같을 때 즉, 해당 숫자의 위에 어떤 숫자도 존재하지 않는 루트일 때까지 재귀함수를 이용해서 최종 루트를 찾는다. 

이 함수를 작성하는 과정에서 TypeError가 발생했다. 

def find(num):
    global root_set
    if num==root_set[num]:
        return num
    else:
        root_set[num]=find(root_set[num])

처음에 find 함수를 위와 같이 작성했었는데, 디버깅을 해보니 union 함수를 실행하는 과정에서 find 함수가 'None'를 반환한 상태에서 root_set 배열 값을 설정해주려다 보니 타입 에러가 발생했다. 즉, root_set[None]에 어떤 값을 할당하거나 root_set[어떤 수]=None을 할당하려 한 것이다.

그 결과, root_set 배열 안에 None값이 들어가다보니 탐색 과정에서 타입에러가 발생하거나 root_set[None]에 어떤 값을 할당하려고 하는 과정에서 타입에러가 발생한 것이다. 

이러한 에러가 발생한 일이 발생한 이유를 찾아보니 

위 예시에선 root_set=[0,1,2,1,4,5,7,7]인 상황에서 find(3)을 할 때 해당 상황이 발생했다. 

위와 같이, return 값이 없는 상태로 find(3)이 종료되는 상황이 발생하기 때문에 생기는 일이었다.

그래서 find 함수를 다음과 같이 수정했다. 

def find(num):
    global root_set
    if num!=root_set[num]:
        root_set[num] = find(root_set[num])
    return root_set[num]

꼼꼼하게 코딩하자...!

그리고 두 집합을 합쳐주는 union 함수는 find 함수를 이용해서 각 수의 루트 노드를 찾아서 같은 루트를 가지도록 할당해준다. 

이러한 과정을 거치다보면 root_set 배열은 각각 자신의 부모 노드를 저장하고 있는 상태가 되고, find 함수를 통해서 최종 루트를 찾아가면서 각 문제에 대한 답을 출력하면 된다. (YES or NO)