[BOJ] #1976 여행가자(Python)

🌱 문제

1976번: 여행 가자

 

🌱 풀이

0. 입력 받기 및 변수 초기화

n = int(input())
m = int(input())
graph = []
parent = [i for i in range(n + 1)]

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

plan = list(map(int, input().split()))

• 전체 도시 개수 n, 여행 도시 개수 m, 도시 연결 정보 graph를 입력받습니다.
• parent는 부모 노드 정보(집합 정보)를 저장할 리스트입니다.

1. Union-Find 알고리즘

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

def union(x, y):
    x = find(x)
    y = find(y)
    if x > y:
        parent[x] = y
    else:
        parent[y] = x

이 문제는 union-find 알고리즘을 사용하여 여행 계획에 있는 도시들이 모두 같은 집합에 있는지 검사하는 문제입니다. 

이때 find함수에서 경로압축을 해주어야 합니다. 경로압축이란, find() 함수에서 부모노드를 찾아가는 과정에서 부모 노드를 업데이트 하는 것을 말합니다. Union Find 알고리즘에 대한 자세한 설명을 아래 링크를 참고하시길 바랍니다.

Union Find 알고리즘 | 개념과 Python 코드 구현

 

2. 집합 만들기

for idx, connected_nodes in enumerate(graph):
    for v in connected_nodes:
        union(idx + 1, v)

 

입력 받은 도시 연결 정보(graph)에 따라 도시가 연결되어있으면 union()하여 parent 리스트를 업데이트해나갑니다.

3. 여행 계획 검사

result = True
root = parent[plan[0]]
for p in plan[1:]:
    if parent[p] != root:
        result = False
        break

if result == True:
    print("YES")
else:
    print("NO")

 

입력받은 여행 계획에서 모든 여행 계획 도시가 같은 부모 노드를 가지면 "YES"를 출력, 그렇지 않으면 "NO"를 출력합니다. 

 

🌱 코드

# !/usr/bin/env python
# -*- coding: utf-8 -*-
# boj 1976 여행가자

n = int(input())
m = int(input())
graph = []
visited = [False] * (n + 1)
parent = [i for i in range(n + 1)]

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

plan = list(map(int, input().split()))

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

def union(x, y):
    x = find(x)
    y = find(y)

    if x > y:
        parent[x] = y
    else:
        parent[y] = x

for idx, connected_nodes in enumerate(graph):
    for v in connected_nodes:
        union(idx + 1, v)

result = True
root = parent[plan[0]]
for p in plan[1:]:
    if parent[p] != root:
        result = False
        break

if result == True:
    print("YES")
else:
    print("NO")