[백준/알고리즘/python/java] 1976번 - 여행 계획

이미지를 누르면 문제 링크를 보실 수 있습니다.

이미지를 누르면 문제 링크를 보실 수 있습니다.

 

이 문제는 유니온-파인드 집합 문제입니다. 

DFS나 BFS로 풀 때, 유의해야 할 점은 결국 자기자신으로 돌아오는 여행루트도 존재한다는 점입니다.

가장 직관적이게 풀 수 있는 알고리즘은 union-find 인데요, 여행계획이 어떻게 주어지더라도 결국 여행이 가능하려면, 각 도시는 같은 사이클에 존재해야 한다는 전제가 있어야 합니다.

하나의 사이클에 같이 존재하지 않는다면, 다시말해 어떠한 도시에서 어떤 도시로 갈 수 있는 방법이 존재하지 않는다면, 여행은 불가능하겠죠.

 

풀이는 코드를 토대로 설명드리겠습니다.

import sys
input = sys.stdin.readline

def find_parent(parent,x):
    if parent[x] != x:
        return find_parent(parent,parent[x])
    return parent[x]

def union_parent(parent,a,b):
    a = find_parent(parent,a)
    b = find_parent(parent,b)
    if a < b :
        parent[b] = a
    else:
        parent[a] = b 
    

N = int(input())    # 도시의 수 N
M = int(input())    # 여행계혹에 속한 도시들의 수 M

graph = [[] for _ in range(N+1)]    # 도시 번호는 1부터니까, 0번째에는 빈 리스트로 초기화
parent = [k for k in range(N+1)]    # 0부터 N까지의 도시의 부모를 자기 자신으로 초기화
# print(parent)
for i in range(1,N+1):
    row = list(map(int,input().strip().split()))
    graph[i].append(row)
 
    for idx,j in enumerate(row):
        if j == 1:  # 연결되어 있다면
            # i와 idx+1 도시 집합 합치기
            union_parent(parent,i,idx+1)  # idx+1(도시번호는 1부터이므로 각 도시번호에 +1 을 해줘야 한다.i는 이미 해당작업 처리)
                                        
plan = list(map(int,input().strip().split()))

cycle = True    # 같은 집합(같은 사이클)인지 판별하는 변수
for idx, city in enumerate(plan[:-1]):
    # 여행계획의 나라의 부모가 같지않으면,(다른집합이라면)
    if (find_parent(parent,city) != find_parent(parent,plan[idx+1])):
        cycle = False   # 다른집합이다.

if cycle == True: # 같은집합이면 여행가능
    print("YES")
else:
    print("NO")

먼저, 도시끼리의 연결정보를 나타내는 graph 리스트와 각 도시의 최종 부모(루트)는 누구인지 나타내는 parent 리스트가 필요합니다.

이 때, 도시는 1번 도시부터인데, 0번째 인덱스를 주의해서 설정해야 합니다. 또, 각 도시의 부모(루트)를 우선 자기자신으로 설정해야 하는 작업도 필요합니다.

 

각 도시마다 0과 1이 포함된 연결정보가 담긴 row를 입력받고, graph에 알맞게 추가합니다. 이 때, row를 돌면서 j 가 1 이라면, (즉 연결된 상태라는 것을 의미하는 1의 값을 만난다면) 해당 도시 i와 row에 담긴 idx+1 도시의 루트를 합쳐줍니다.

서로 연결되어 있으니까 루트는 같으니까요!

 

여기서 대체 루트가 뭔데? 하실 수 있으실겁니다! 루트는 말 그대로 어떠한 노드끼리 연결되었을 때 가장 작은 번호를 나타내는 노드입니다.

예를 들어, 노드 1,노드 2, 노드 3( 뭐, 여기서는 도시겠죠..?) 이 연결되어 있다면, 이 사이클(집합)의 루트노드는 노드1이 될 것입니다. 

 

이런 작업이 union_parent함수에 구현되어 있습니다.

def find_parent(parent,x):
    if parent[x] != x:
        return find_parent(parent,parent[x])
    return parent[x]

def union_parent(parent,a,b):
    a = find_parent(parent,a)
    b = find_parent(parent,b)
    if a < b :
        parent[b] = a
    else:
        parent[a] = b

union_parent함수를 살펴보면, find_parent함수를 먼저 살펴봐야 할 것입니다.

find_parent함수는 말그대로 부모(루트)를 찾는 로직입니다. 재귀를 통해서 그 사이클의 부모를 계속해서 찾아나가는 과정이 되겠죠. 

다시 union_parent 함수로 돌아온다면, a노드와 b노드의 부모를 찾은 후에, a노드의 부모가 b노드의 부모보다 더 작은 수라면, b노드의 부모를 a노드의 부모로 갱신합니다.(더 작은 수가 부모노드로 설정되도록 하기 위한 작업)

반대라면, a노드의 부모를 b노드의 부모로 갱신하겠죠?

 

자, 거의 다 왔습니다.

이제 입력받은 여행계획을 순회하면서, 그 다음으로 여행할 도시의 부모(루트)와 지금 순회하고 있는 도시의 부모(루트)가 다르다면, 두 도시는 서로 다른 집합에 있는 것입니다. 즉, 여행할 수 없는 계획이라는 것이죠. cycle이 False가 된다면 남은 여행계획을 볼 필요 없이 바로 break해서 나와도 됩니다. (위 코드에서는 break는 안했습니다!)

 

여행이 가능하다면, YES를 출력해주고

두 도시가 "서로 다른 집합, 즉 다른 사이클에 있으면, 다시말해 부모(루트)가 다른 도시라면!!", 여행할 수 없는 계획이기 때문에, NO를 출력해주면 됩니다!

 

플로이드 워셜 알고리즘으로도 풀 수 있지만, 유니온-파인드 알고리즘이 조금 더 시간복잡도가 낮기 때문에 이런 로직으로 정리해봤습니다.