[BOJ] 1976 : 여행가자 [유니온파인드](Python)

문제 링크

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

1976번: 여행 가자

소스 코드

n = int(input())
m = int(input())

graph = []
for i in range(n):
    graph.append(list(map(int,input().split())))

plan = list(map(int,input().split()))
parent = [x for x in range(n+1)] # 0부터 n까지

def find_parent(x):
    if parent[x] != x: # 재귀적으로 부모 노드 탐색
        parent[x] = find_parent(parent[x])
    return parent[x] # 최상단 부모 노드 반환

def union_parent(a,b):
    a = find_parent(a)
    b = find_parent(b)
    if a < b: # 작은 숫자가 큰 숫자의 부모
        parent[b] = a
    else:
        parent[a] = b

for i in range(n-1): # 그래프 우상단 삼각형만 탐색
    for j in range(i+1,n):
        if graph[i][j] == 1:
            union_parent(i+1,j+1)

for i in range(len(plan)-1): # 부모가 동일하면 continue
    if find_parent(plan[i]) == find_parent(plan[i+1]):
        continue
    else: # 부모가 다른 경우 존재하면 프로그램 종료
        print('NO')
        exit()
print('YES')

 

해설

• 그래프 탐색, 유니온 파인드, 서로소 집합

예전에 알고리즘을 공부할 때 이런 유형들을 접하면서 알고리즘도 개념을 배우지 않고서는 혼자 풀어낼 수 없겠다고 생각했었습니다.

이 문제는 서로소 집합/유니온 파인드 유형에 속합니다.

노드 간의 연결 관계를 통해 노드 간의 부모/자식 관계를 파악하는 유형이죠.

 

문제에서 주어진 여행 계획이 가능한 것인지 그렇지 않은지를 판단하는 기준은 여행 계획에 주어진 노드들이 동일한 부모 노드를 갖는가입니다.

만약 주어진 노드들의 부모 노드가 모두 동일하다면 하나의 연결 그래프를 형성한다는 뜻이기 때문이죠.

개념에 대한 자세한 내용은 이코테 강의를 보면서 이전에 네이버 블로그에 포스팅한 것을 참고하셔도 좋을 것 같습니다.

 

부모 노드 찾기

def find_parent(x):
    if parent[x] != x: # 재귀적으로 부모 노드 탐색
        parent[x] = find_parent(parent[x])
    return parent[x] # 최상단 부모 노드 반환

parent 라는 리스트는 인덱스 자체의 값으로 초기화되어 있습니다.

그러나 탐색 과정에서 부모 노드가 업데이트 되어 있다면 조건문의 parent[x] != x 에 해당할 것입니다.

예를 들어 [0, 1, 2, 3] 으로 초기화되어 있던 parent가 [0, 1, 2, 1] 로 변경된 상태라면 parent[3] != 3 을 만족하는 것입니다.

이때는 부모노드의 부모노드가 추가로 존재하는지 확인하여 최상단 부모 노드가 반환될 수 있도록 짠 코드입니다.

 

부모 노드 업데이트하기

def union_parent(a,b):
    a = find_parent(a)
    b = find_parent(b)
    if a < b: # 작은 숫자가 큰 숫자의 부모
        parent[b] = a
    else:
        parent[a] = b

두 노드 a, b 가 연결되어 있을 경우 위 함수에 인자로 전달하여 부모 노드를 업데이트 합니다.

먼저 각 노드의 부모 노드를 찾습니다.

그리고 부모 노드의 대소를 비교하여 그 값이 더 작은 것이 큰 것의 부모 노드로 업데이트 됩니다.

 

위 예시를 다시 생각해보면 1번 노드와 3번 노드가 연결되어 있을 경우 1번 노드가 3번 노드의 부모가 되는 것입니다.

이렇게 하는 이유는 일반적으로 오름차순이 가정되어 있기 때문이고, 문제의 조건에 따라 이는 달라질 수 있습니다.

 

부모 노드 비교하기

for i in range(len(plan)-1): # 부모가 동일하면 continue
    if find_parent(plan[i]) == find_parent(plan[i+1]):
        continue
    else: # 부모가 다른 경우 존재하면 프로그램 종료
        print('NO')
        exit()
print('YES')

부모 노드가 무엇인지 탐색, 비교하여 업데이트하는 과정이 모두 끝났다면 마지막으로는 plan에 담긴 노드들의 부모가 동일한지 확인해줍니다.

만약 동일하다면 문제가 없는 것이므로 continue를 통해 반복문을 다 돌아도 아무 문제 없이 'YES'가 출력되도록 합니다.

그러나 단 한 번의 경우에서라도 부모 노드가 다른 케이스가 발견되면 'NO'를 출력하고 프로그램을 종료합니다.