[백준] 1976번 여행 가자 파이썬 풀이

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

1976번: 여행 가자

---

1. 접근 방식 

• 도시 번호가 1부터 N까지이지만, 편의를 위해 0부터 N - 1까지로 변환하자.

 

도시 정보를 담은 2차원 리스트(maps)의 행과 열의 인덱스는 0부터 N - 1까지이다. 따라서 개발의 편의를 위해 도시 번호를 0부터 시작하게 변경하자.

다만 마지막에 여행 계획이 담긴 도시들의 번호는 문제에서 정한 대로 번호가 1부터 시작하므로, 그 번호에 1을 감소시킨 도시 번호로 연결 정보를 확인해야 한다. 예를 들어 도시 1, 2, 3을 방문할 수 있는지 확인하기 위해 도시 번호를 0, 1, 2로 치환하여야 한다.

 

 

• 여러 도시가 존재하고, 그 도시가 서로 연결되어 있는지 확인한다 ➡️ union-find 알고리즘

도시 정보를 담은 2차원 리스트(maps)를 순회하면서 도시의 연결 정보를 parents 리스트에 담는다. 위 그림에서 왼쪽은 maps이고, 오른쪽은 도시 연결 정보가 담긴 parents이다. 도시 0과 1을 연결하면 도시 0과 1은 서로 같은 그룹에 속한다는 걸 표현하고 있다.

 

 

 

2. 정답 코드

import sys
sys.setrecursionlimit(10 ** 6)
input = sys.stdin.readline

def find(parents, x):
    if parents[x] == x:
        return x

    parents[x] = find(parents, parents[x])
    return parents[x]

def union(parents, a, b):
    a = find(parents, a)
    b = find(parents, b)

    if a > b:
        parents[a] = b
    else:
        parents[b] = a

def is_connected(parents, cities):
    pre = find(parents, cities[0])

    for city in cities[1:]:
        cur = find(parents, city)
        if pre != cur:
            return "NO"
        pre = cur

    return "YES"

def main():
    n = int(input())
    m = int(input())
    maps = [list(map(int, input().rstrip().split())) for _ in range(n)]
    cities = list(map(lambda x: int(x) - 1, input().rstrip().split()))

    parents = [i for i in range(n)]
    for i in range(n):
        for j in range(n):
            if maps[i][j] == 1:
                union(parents, i, j)
    print(is_connected(parents, cities))

main()