[BOJ / 백준 1976번][PYTHON] 여행 가자 (유니온 파인드)

유니온 파인드를 연습하기 위한 두 번째 문제이다. 

입력의 의미는 도시가 어느 도시와 연결되었는지에 대한 여부이다. 원하는 여행 경로에 따라 해당 여행 경로가 가능한 여행 경로면 YES아니면 NO가 출력되는 형식이다.

 

해당 문제를 보며 처음 든 생각은 이거를 만약 코딩 테스트나 문제 구분 없이 푼다면 당연히 BFS로 풀 거 같다는 생각을 하였다. (물론 지금도 마찬가지이다.) 하지만 분류를 유니온 파인드에서 발견하였으므로 해당 방법으로 풀겠다. 유니온 파인드에 대한 구현은 그 전(골드 5 [1717번] - 집합의 표현)과 같게 생각했다. 알고리즘 자체가 어려운 것도 아니고 굉장히 간단했기 때문에 기본으로 이 함수를 적어두고 어떻게 해결할 까 생각했다. 유니온 파인드의 구현 방식을 생각해보면 같은 집합에 속하는 수는 같은 수를 value로 갖게 된다. bfs에서 같은 큐에 들어가는 상황 자체가 마찬가지로 같은 value로 들어가는 것과 동일하게 여기면 되었다. 구현은 아래와 같다.

import sys

n = int(sys.stdin.readline())
m = int(sys.stdin.readline())

arr = [i for i in range(n+1)]

def find(x):
    if arr[x] != x:
        arr[x] = find(arr[x])
    return arr[x]

def union(a,b):
    a = find(a)
    b = find(b)
    if a < b:
        arr[b] = a
    else:
        arr[a] = b

for i in range(1, n+1):
    tmp = list(map(int, sys.stdin.readline().split()))
    for j in range(1, n+1):
        if tmp[j-1] == 1:
            union(i,j)

result = list(map(int, sys.stdin.readline().split()))
compare = find(result[0])
flag = True

for i in result:
    if find(i) != compare:
        flag = False
        break

if flag == False:
    print("NO")
else:
    print("YES")

입력이 들어오게 되는데, 1번 줄에는 1번 (A 마을)과 연결된 타 마을이 0(연결 x) 1(연결 O) 중 하나로 입력이 들어오게 된다. 그리하여 입력이 들어오는 경우 1이라면 해당 마을과 Union 연산을 진행하게 된다. 마지막 입력으로는 여행 경로가 입력되게 된다. 나는 이 부분을 result라는 이름의 변수로 저장하였다. 그 후 result 배열의 값 하나하나를 꺼내서 find(i)로 비교연산을 하게 된다.

만약 find(result[0]) = 출발지와 같은 집합에 속해있다면 통과지만, 같은 집합에 속해 있지 않다면 해당 여행 경로는 불가능 하기 때문에 NO를 출력해야 할 것이다. if문에 해당하는 경우 print("NO")와 exit(0)을 통하여 바로 할 수도 있지만, 오늘따라 그냥 flag가 쓰고 싶었다. 뭔가 말하는 것처럼 코딩하는 느낌이라던지 그런 느낌이었다. 큰 이유는 없다. 자기 스타일대로 하면 될 거 같다.

 

위와 같이 코딩을 진행하여 결과는 ?

장렬하게 맞았다.

 

성공했당