[Swift] 프로그래머스 주사위 고르기

프로그래머스

문제

• A와 B가 n개의 육면체 주사위로 승부
• 각 면이 나올 확률이 동일
• A와 B는 n / 2 개씩 가져가 주사위를 굴린 뒤, 나온 수를 합한 점수로 계산
• A는 주어진 주사위에서 승리할 확률이 가장 높은 주사위를 선택

 

제한사항

• 승리할 확률이 가장 높은 주사위 조합으로 주어짐
• n은 2의 배수
• dice의 개수는 10개 이하
• dice에 적힌 수는 100이하 자연수

 

풀이

1. 각 주사위를 고를 수 있는 조합(combination) 구하기
2. n/2개를 던져서 나올 수 있는 조합 구하기
3. A가 고른 주사위들(1번의 원소들)을 2번의 조합을 통해 점수를 합산
4. B가 고른 주사위들의 점수 합산과 3번을 비교
5. 가장 확률이 높은 주사위들을 반환

큰 흐름을 보게되면 아래 코드와 같다

// 1번: 선택된 주사위들의 조합
let selectedDices = selectDicesCombination()
// 2번: 굴린 주사위들 조합
let rollings = rollingDiceCombination()

// 정답
var result: [Int] = []
var win = 0

// 3번: 1번에서 하나씩 확인
for selectedDice in selectedDices {
    var A = [Int](), B = [Int]()
    selectedDice.enumerated().forEach { i, v in
        v ? (A.append(i)) : (B.append(i))
    }
    
    // 3 ~ 4번: 주사위들 결과값 합산
    let APoints = sumDices(rollings, A)
    let BPoints = sumDices(rollings, B)

    let winInfo = countWin(APoints, BPoints)

    if win < winInfo.0 {
        win = winInfo.0
        result = A
    }
    if win < winInfo.1 {
        win = winInfo.1
        result = B
    }
}

return result.map { $0 + 1 }

 

Combination과 다른 점은 A가 선택하면 B는 정반대의 주사위를 선택하기 때문에

선택지를 절반으로 줄이는 효과를 가질 수 있다.

 

선택된 주사위들을 구하는 조합

반복되는 구간을 1부터 시작하게 되면 조합의 절반을 선택하는 효과를 보여준다.

// 주사위를 뽑는 경우의 수
func selectDicesCombination() -> [[Bool]] {
    var ret = [[Bool]]()
    var visit = Array(repeating: false, count: n)
    func dfs(_ i: Int, _ c: Int) {
        if c == n / 2 { ret.append(visit); return }

        for i in i..<n {
            visit[i] = true
            dfs(i+1, c+1)
            visit[i] = false
        }
    }
    dfs(1, 0)
    return ret
}

 

예를 들어, 주사위가 4개를 던진다면 아래와 같은 결과를 반환한다.

[[false, true, true, false], [false, true, false, true], [false, false, true, true]]

// (2, 3) (1, 4)
// (2, 4) (1, 3)
// (3, 4) (1, 2)

 

n/2 주사위를 던져서 나오는 조합

모든 주사위는 육면체이기 때문에, 6^(n/2) 의 조합 수가 나온다.

// 주사위를 굴리는 모든 경우의 수
func RollingDiceCombination() -> [[Int]] {
    var ret = [[Int]]()
    func dfs(_ arr: [Int]) {
        if arr.count == n / 2 { ret.append(arr); return }

        var arr = arr
        for i in 0..<6 {
            arr.append(i)
            dfs(arr)
            arr.removeLast()
        }
    }
    dfs([])

    return ret
}

 

선택된 주사위에서 나올 수 있는 합산들 그리고 승률 확인

이분탐색을 이용하여 승률 확인

// 각 조합에 따른 합산
func sumDices(_ picks: [[Int]], _ dices: [Int]) -> [Int] {
    var ret = [Int]()
    for pick in picks {
        var value = 0
        for (p, d) in zip(pick, dices) {
            value += dice[d][p]
        }
        ret.append(value)
    }
    return ret
}


// 합산을 토대로 이분탐색으로 확인
func countWin(_ A: [Int], _ B: [Int]) -> Int {
    var result = 0, A = A, B = B
    A.sort()
    B.sort()

    for a in A {
        var start = 0, end = B.count - 1
        while start <= end {
            let mid = (start + end) / 2
            if a > B[mid] {
                start = mid + 1
            } else {
                end = mid - 1
            }
        }
        result += start
    }

    return result
}

 

하지만 Swift에서는 시간을 초과가 나게 된다.

합산된 값들을 확인하니 겹치는 숫자들이 많기 때문에 Dictionary로 표현하니 해결됐다.

함수 countWin은 A와 B의 승리하는 횟수를 튜플로 반환한다.

수정된 코드는 아래와 같다.

func sumDices(_ picks: [[Int]], _ dices: [Int]) -> [Int: Int] {
    var ret = [Int: Int]()
    for pick in picks {
        var value = 0
        for (p, d) in zip(pick, dices) {
            value += dice[d][p]
        }
        ret[value, default: 0] += 1
    }    
    return ret
}

func countWin(_ A: [Int: Int], _ B: [Int: Int]) -> (Int, Int) {
    var result = (0, 0)

    for aKey in A.keys {
        for bKey in B.keys {
            if aKey > bKey {
                result.0 += A[aKey]! * B[bKey]!
            } else if aKey < bKey {
                result.1 += A[aKey]! * B[bKey]!
            }
        }
    }

    return result
}

 

카카오 문제를 보게되면 문제를 잘 읽고 구현만 하면 된다.
시간 초과가 발생하겠지만... 이분 탐색, Dictionary, Queue 등 다양한 방법을 빠르게 해결할 줄 알아야 할 거 같은 문제였다. 나름 고민을 했는데도 시간 초과할 줄이야

 

코드

 

import Foundation

func solution(_ dice:[[Int]]) -> [Int] {
    let n = dice.count

    // 주사위를 뽑는 경우의 수
    func selectDicesCombination() -> [[Bool]] {
        var ret = [[Bool]]()
        var visit = Array(repeating: false, count: n)
        func dfs(_ i: Int, _ c: Int) {
            if c == n / 2 { ret.append(visit); return }

            for i in i..<n {
                visit[i] = true
                dfs(i+1, c+1)
                visit[i] = false
            }
        }
        dfs(1, 0)
        return ret
    }

    // 주사위를 굴리는 모든 경우의 수
    func rollingDiceCombination() -> [[Int]] {
        var ret = [[Int]]()
        func dfs(_ arr: [Int]) {
            if arr.count == n / 2 { ret.append(arr); return }

            var arr = arr
            for i in 0..<6 {
                arr.append(i)
                dfs(arr)
                arr.removeLast()
            }
        }
        dfs([])

        return ret
    }

    // 각 조합에 따른 합산
    func sumDices(_ picks: [[Int]], _ dices: [Int]) -> [Int: Int] {
        var ret = [Int: Int]()
        for pick in picks {
            var value = 0
            for (p, d) in zip(pick, dices) {
                value += dice[d][p]
            }
            ret[value, default: 0] += 1
        }
        return ret
    }

    func countWin(_ A: [Int: Int], _ B: [Int: Int]) -> (Int, Int) {
        var result = (0, 0)

        for aKey in A.keys {
            for bKey in B.keys {
                if aKey > bKey {
                    result.0 += A[aKey]! * B[bKey]!
                } else if aKey < bKey {
                    result.1 += A[aKey]! * B[bKey]!
                }
            }
        }

        return result
    }
    // 1번: 선택된 주사위들의 조합
    let selectedDices = selectDicesCombination()
    // 2번: 굴린 주사위들 조합
    let rollings = rollingDiceCombination()

    // 정답
    var result: [Int] = []
    var win = 0

    // 3번: 1번에서 하나씩 확인
    for selectedDice in selectedDices {
        var A = [Int](), B = [Int]()
        selectedDice.enumerated().forEach { i, v in
            v ? (A.append(i)) : (B.append(i))
        }

        // 3 ~ 4번: 주사위들 결과값 합산
        let APoints = sumDices(rollings, A)
        let BPoints = sumDices(rollings, B)

        let winInfo = countWin(APoints, BPoints)

        if win < winInfo.0 {
            win = winInfo.0
            result = A
        }
        if win < winInfo.1 {
            win = winInfo.1
            result = B
        }
    }

    return result.map { $0 + 1 }
}