[코딩테스트] 🗺️ 여행경로 - 백준 1976

문제링크

난이도

골드 4

유형

Union-Find

소요시간

1시간

회고

유니온 파인드를 저번 표병합 때 대충 넘겼더니 머리속에서 사라졌다. 이번에 개념을 제대로 잡는 계기가 되었다.

코드

let N = Int(readLine()!)!; let M = Int(readLine()!)!
var answer = "YES"
var group = Array(0...N)

// 그룹의 대표 찾기
func find(_ x: Int) -> Int {
    if group[x] == x {
        return group[x]
    }

    group[x] = find(group[x])
    return group[x]
}

// 대표가 다르면 단일화
func union(_ x: Int, _ y: Int) {
    let rootX = find(x)
    let rootY = find(y)

    if rootX == rootY {
        return
    }

    group[rootY] = rootX
}

(1...N).forEach { i in
    let row = readLine()!.split(separator: " ").compactMap { Int($0) }
    (1...N).forEach { j in
        if row[j - 1] == 1 {
            // 단일화
            union(i, j)
        }
    }
}

let path = readLine()!.split(separator: " ").compactMap { Int($0) }

let start = path[0]
for step in (0..<path.count) {
    if find(start) != find(path[step]) {
        answer = "NO"
        break
    }
}

print(answer)