[Algorithm] 백준 - 14891번 톱니바퀴 (Swift)

문제 소개

https://www.acmicpc.net/problem/14891
엄청난 구현문제... (?)
 

문제 풀이

톱니바퀴 돌리기

큐로 구현하는것이 일반적이지만 여기서는 배열의 갯수와 성분들이 작아서 간단하게 removeLast, removeFirst 메서드로 구현했다.

func rotateGear(_ i: Int, _ isClockwise: Bool) {
    if isClockwise {
        gears[i].insert(gears[i].removeLast(), at: gears[i].startIndex)
    } else {
        gears[i].append(gears[i].removeFirst())
    }
}

 

같이 돌아가는 톱니바퀴 그룹화

union, find 알고리즘을 사용했다. 결국 같은 그룹안에 있어야 톱니바퀴가 돌아가는 것이기 때문에 그룹을 짓고 parents 를 반환하는 함수를 만들었다.

func groupGears(_ n: Int) -> [Int] {
    var parents = [0, 1, 2, 3]
    
    func find(_ a: Int) -> Int {
        return parents[a] == a ? a : find(parents[a])
    }
    
    func union(_ a: Int, _ b: Int) {
        let a = find(a)
        let b = find(b)
        parents[b] = a
    }
    
    for i in 0..<3 {
        if gears[i].prefix(3).last != gears[i+1].prefix(7).last {
            union(i, i+1)
        }
    }
    
    return parents
}

 

점수 계산

이 부분은 단순하게 문제에서 주어진대로 구현했다.

func getScore() -> Int {
    var ans = 0
    
    for i in gears.indices {
        if gears[i].first == "1" {
            ans += Int(pow(2.0, Double(i)))
        }
    }
    
    return ans
}

 

정답 코드

import Foundation

var gears = [String]()

for _ in 0..<4 {
    let gear = readLine()!
    gears.append(gear)
}

let k = Int(readLine()!)!

// 같이 돌아가는 톱니바퀴 그룹화
func groupGears(_ n: Int) -> [Int] {
    var parents = [0, 1, 2, 3]
    
    func find(_ a: Int) -> Int {
        return parents[a] == a ? a : find(parents[a])
    }
    
    func union(_ a: Int, _ b: Int) {
        let a = find(a)
        let b = find(b)
        parents[b] = a
    }
    
    for i in 0..<3 {
        if gears[i].prefix(3).last != gears[i+1].prefix(7).last {
            union(i, i+1)
        }
    }
    
    return parents
}

// 톱니바퀴 돌리기
func rotateGear(_ i: Int, _ isClockwise: Bool) {
    if isClockwise {
        gears[i].insert(gears[i].removeLast(), at: gears[i].startIndex)
    } else {
        gears[i].append(gears[i].removeFirst())
    }
}

// 점수 구하기
func getScore() -> Int {
    var ans = 0
    
    for i in gears.indices {
        if gears[i].first == "1" {
            ans += Int(pow(2.0, Double(i)))
        }
    }
    
    return ans
}

for _ in 0..<k {
    let nd = readLine()!.components(separatedBy: " ").map { Int($0)! }
    let n = nd[0] - 1
    let d = nd[1] == 1 ? true : false // 시계방향 true, 반시계방향 false
    
    let parents = groupGears(n)
        
    for i in 0..<4 {
        if parents[n] == parents[i] {
            if abs(n-i) % 2 == 0 {
                rotateGear(i, d)
            } else {
                rotateGear(i, !d)
            }
        }
    }
}

print(getScore())

 

디버깅

아래와 같이 gears 를 입력받으면 모양을 출력해주는 함수를 임의로 구현하여 검증했다.

func printGears(_ gears: [String]) {
    func get(_ i: Int, _ j: Int) -> String {
        return String(gears[i].prefix(j+1).last ?? "?")
    }
    
    print("  \(get(0, 0))     \(get(1, 0))     \(get(2, 0))     \(get(3, 0))  ")
    print(" \(get(0, 7)) \(get(0, 1))   \(get(1, 7)) \(get(1, 1))   \(get(2, 7)) \(get(2, 1))   \(get(3, 7)) \(get(3, 1)) ")
    print("\(get(0, 6))   \(get(0, 2)) \(get(1, 6))   \(get(1, 2)) \(get(2, 6))   \(get(2, 2)) \(get(3, 6))   \(get(3, 2))")
    print(" \(get(0, 5)) \(get(0, 3))   \(get(1, 5)) \(get(1, 3))   \(get(2, 5)) \(get(2, 3))   \(get(3, 5)) \(get(3, 3)) ")
    print("  \(get(0, 4))     \(get(1, 4))     \(get(2, 4))     \(get(3, 4))  ")
}

printGears(gears)

// 출력 값
  1     1     1     0  
 1 1   0 1   0 1   1 0 
1   0 1   1 1   0 0   0
 1 1   0 1   1 0   0 0 
  0     1     1     0