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

https://school.programmers.co.kr/learn/courses/30/lessons/258709

프로그래머스

 

2024.01.22 기준 Level 3

 

알고리즘 공부를 위해 풀고 기록하는 글입니다!
참고만 해주시고 더 좋은 풀이법이 있다면 알려주세요!

 

아이디어

문제 조건

• n개의 주사위 존재 (2<=n<=10)
• 각 주사위는 6개의 수를 가진다.
• A와 B는 각자 이 n개의 주사위에서 n/2개의 주사위를 가져간다. (서로 겹치지 않게 가져감)
• 이 때 각자 가져간 주사위를 모두 굴린 뒤, 나온 수들을 모두 합해 점수를 계산한다.
• 점수가 더 큰 쪽이 승리하며 점수가 같다면 무승부이다.
• A는 자신이 가장 승리할 확률이 높아지도록 주사위를 가져가려고 한다.
• 주사위에 적힌 수를 담은 배열인 dice가 주어질 때 A의 승률를 최대로 하는 주사위 경우의 수를 리턴하자

 

우선 2 <= n <= 10 이라는 점이 중요합니다.

주사위가 최대 10개까지이기 때문에 큰 수가 아닙니다.

즉, 모든 경우의 수를 완전탐색해서 A의 승률이 최대인 경우를 찾아야 하는 문제였습니다.

 

하지만 시간 복잡도를 생각하지 않고 무작정 모든 경우의 수를 찾을 경우 시간초과가 발생합니다.

n이 10이라고 하면,

A와 B가 주사위를 나누어 가지는 경우의 수는 10C5 = 252 이며,

이 각각의 경우의 수에서 주사위를 굴렸을 때 나오는 경우의 수는 A에서 6^5개, B에서 6^5개입니다.

A와 B가 같이 주사위를 굴리기 때문에 전체 경우의 수는 (6^5) * (6^5) = 6^10 개 입니다.

이는 매우 큰 연산 횟수이기 때문에 시간 초과가 발생합니다.

 

따라서 완전 탐색을 하되 시간 복잡도를 최대한 낮춰야 합니다.

 

➡️ A와 B의 주사위를 한꺼번에 시뮬레이션해서 승패를 계산하는 것이 아니라, A와 B를 따로 시뮬레이션해야 합니다.

 

먼저 A의 주사위를 굴린 결과를 시뮬레이션해 6^(n/2) 가지의 주사위 눈의 합을 배열(diceSumForA)에 저장합니다.

n이 10일 경우 최대 6^5 = 7776 개가 생깁니다.

 

B도 마찬가지로 시뮬레이션해 눈의 합을 배열(diceSumForB)에 저장합니다.

 

그리고 diceSumForA의 원소를 반복문으로 돌며 해당 값보다 diceSumForB에서 작은 값들의 수를 세어 그 값들을 모두 합하면 A가 승리하는 경우의 수와 같아집니다.

 

물론 이걸 2중 for문으로 계산하게 되면 최악의 경우 7776 * 7776 번의 연산이 필요하기 때문에 알고리즘 기법을 사용해야 합니다.

 

저는 이분탐색을 사용했습니다.

diceSumForB를 오름 차순으로 정렬하고 diceSumForA의 원소 num보다 큰 값이 시작되는 인덱스를 구하여 그보다 작은 인덱스들은 num보다 작은 수가 됩니다.

 

이렇게 하면 diceSumForB의 길이를 M이라고 했을 때 O(M*LogM)으로 연산을 수행할 수 있습니다.

 

 

풀이 순서를 다시 정리하면 다음과 같습니다.

1. A와 B가 나누어 가질 주사위 경우의 수를 구합니다. (조합 사용)
2. 각자가 가진 주사위를 굴렸을 때 나올 수 있는 눈의 합을 배열에 저장합니다.
3. 2번에서 구한 서로의 배열에서 A가 이기는 경우의 수의 개수를 구합니다. (이분 탐색 사용)
4. 3번에서 구한 개수가 최대가 될 때를 리턴합니다.

 

 

 

 

풀이

1. combination 함수 구현

fileprivate func combination<T: Comparable>(_ array: [T], _ n: Int) -> [[T]] {
    var result = [[T]]()
    if array.count < n { return result }
    
    var stack = array.enumerated().map { ([$0.element], $0.offset) }
    
    while stack.count > 0 {
        let now = stack.removeLast()
        
        let elements = now.0
        let index = now.1
        
        if elements.count == n {
            result.append(elements.sorted())
            continue
        }
        
        guard index+1 <= array.count-1 else { continue }
        
        for i in index+1...array.count-1 {
            stack.append((elements + [array[i]], i))
        }
    }
    
    return result
}

 

A와 B가 나누어 가질 주사위의 경우 수를 구하는 조합 함수입니다.

 

 

2. getAllDiceSum 함수 구현

fileprivate func getAllDiceSum(dice: [[Int]]) -> [Int] {
    var res = [Int]()
    
    func recursion(index: Int, sum: Int) {
        if index == dice.count {
            res.append(sum)
            return
        }
        
        for i in 0..<6 {
            recursion(index: index+1, sum: sum + dice[index][i])
        }
    }
    
    recursion(index: 0, sum: 0)
    
    return res
}

 

자신이 가진 주사위들을 굴렸을 때 나올 수 있는 눈의 합의 모든 경우의 수를 구하는 함수입니다.

재귀를 사용했습니다.

 

3. calculateAWinCount 함수 구현

fileprivate func calculateAWinCount(diceSumForA: [Int], diceSumForB: [Int]) -> Int {
    let diceSumForB = diceSumForB.sorted()
    
    var winCnt = 0
    
    for num in diceSumForA {
        // 이분 탐색
        var lo = 0
        var hi = diceSumForB.count
        
        while lo + 1 < hi {
            let mid = (lo + hi) / 2
            let midNumInB = diceSumForB[mid]
            
            if num <= midNumInB {
                hi = mid
            } else if num > midNumInB {
                lo = mid
            }
        }
        
        winCnt += lo + 1
    }
    
    return winCnt
}

 

A가 주사위를 굴렸을 때 나온 눈의 합 배열과 B가 주사위를 굴렸을 때 나온 눈의 합 배열을 파라미터로 받아서 A가 이기는 횟수를 리턴하는 함수입니다.

 

이분 탐색을 사용하기 위해 diceSumForB를 오름차순으로 정렬하고 시작합니다.

 

 

4. solution 함수 구현

fileprivate func solution(_ dice:[[Int]]) -> [Int] {
    let count = dice.count
    let combinationResults = combination((0..<count).map { $0 }, count/2)
    
    var maxWinCnt = 0
    var result = [Int]()
    
    for diceForA in combinationResults {
        let diceForB = (0..<count).filter { !diceForA.contains($0) }
        let diceSumForA = getAllDiceSum(dice: diceForA.map { dice[$0] })
        let diceSumForB = getAllDiceSum(dice: diceForB.map { dice[$0] })
        
        var winCnt = calculateAWinCount(diceSumForA: diceSumForA, diceSumForB: diceSumForB)
        
        if winCnt > maxWinCnt {
            maxWinCnt = winCnt
            result = diceForA
        }
    }
    
    return result.map { $0+1 }
}

 

앞서 구현한 함수들을 차례로 호출하여 문제를 해결합니다.

 

combination 함수를 호출하여 A와 B가 주사위를 나누어 가진 모든 경우의 수를 구하고,

getAllDiceSum 함수를 호출하여 각 경우의 수에서 주사위를 굴렸을 때 나올 수 있는 눈의 합을 구합니다.

calculateAWinCount 함수를 호출하여 A가 이기는 횟수를 구하고,

이 횟수가 현재까지 최대라면 result를 갱신합니다.

 

 

마지막에는 주사위의 번호를 리턴해야 하기 때문에 reuslt의 원소들을 1씩 더해서 리턴합니다.

 

 

마무리

실제로 카카오 인턴 코딩테스트를 볼 때는 완전 탐색을 해야 하는 것은 알고 있었지만 시간 초과가 발생하는 로직만 생각이 나서 모든 테스트 케이스를 통과하지 못했던 문제였습니다.

 

침착하게 잘 생각해 보면 A와 B를 동시에 시뮬레이션 할 필요 없이 각자 시뮬레이션하고 두 결과를 비교할 때는 이분탐색으로 시간 복잡도를 대폭 낮출 수 있었습니다.

 

역시 차분하게 단계별로 문제 풀이 아이디어를 찾아가는 과정이 중요한 것 같습니다.