[파이썬, 스위프트][백준 1976][💛골드 4] 여행 가자 (유니온 파인드)

🧞‍♂️ 난이도

💛 골드 4

🧞‍♂️ 문제

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

🧞‍♂️ 풀이 방법

이 문제는 예시 테스트 케이스 그려보면 바로 유니온 파인드 알고리즘이 떠오를 수 있습니다.

 4 - 0 - 1 - 2
    \  /
     3

여기서 여행 경로가 4-2-1-2-3 이면

4에서 2 까지 경로가 있는지

2에서 1까지 경로가 있는지

1에서 2까지 경로가 있는지

2에서 3까지 경로가 있는지

전부 확인해보면 되는 문제인데, 위 그림은 한눈에 봐도 모든 경로가 존재하다는걸 알 수 있죠

경로가 없는 경우는 다음과 같은 경우죠

 4 - 0 - 1 - 2
 
     3

즉, 모든 원소가 하나의 집합에 속해있는지 확인하면 됩니다.

유니온 파인드로 모든 원소의 부모를 찾은 후에, 주어진 여행 경로의 원소들이 전부 부모가 같은지만 확인하면 됩니다.

🧞‍♂️ 파이썬 코드

import sys
input = sys.stdin.readline

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

def union(parent, a, b):
    aParent = find(parent, a)
    bParent = find(parent, b)
    if aParent < bParent:
        parent[bParent] = aParent
    else:
        parent[aParent] = bParent

N = int(input())
M = int(input())

parents = [i for i in range(N)]

for i in range(N):
    data = map(int, input().split())

    for j, isPath in enumerate(data):
        if isPath:
            union(parents, i, j)
            
travelPath = list(map(int, input().split()))

for i in range(1, len(travelPath)):
    start = travelPath[i-1] - 1
    destination = travelPath[i] - 1

    if find(parents, start) != find(parents, destination):
        print("NO")
        exit(0)
print("YES")

🧞‍♂️ 스위프트 코드

import Foundation

func find(_ parent: inout [Int], _ node: Int) -> Int {
    if parent[node] != node {
        parent[node] = find(&parent, parent[node])
    }
    return parent[node]
}

func union(_ parent: inout [Int], _ a: Int, _ b: Int) {
    let aParent = find(&parent, a)
    let bParent = find(&parent, b)

    if aParent < bParent {
        parent[bParent] = aParent
    } else {
        parent[aParent] = bParent
    }
}

let N = Int(readLine()!)!
let M = Int(readLine()!)!

var parent: [Int] = []

for i in 0..<N {
    parent.append(i)
}

for i in 0..<N {
    let data = readLine()!.split(separator: " ").compactMap { Int($0) }

    for (j, isPath) in data.enumerated() {
        if isPath == 1 {
            union(&parent, i, j)
        }
    }
}
let travelPath = readLine()!.split(separator: " ").compactMap { Int($0)! - 1 }

for i in 1..<travelPath.count {
    let start = travelPath[i-1]
    let destination = travelPath[i]
    
    if find(&parent, start) != find(&parent, destination) {
        print("NO")
        exit(0)
    }
}
print("YES")

---

맨 위로 이동하기

BOJ 카테고리 내 다른 글 보러가기