[프로그래머스 / python / Level 3] 주사위 고르기

문제

 

2024 KAKAO WINTER INTERNSHIP 3번 문제이다.

 

아주 어려운 문제였다. 알고리즘 분류도 애매하다. 완전탐색, 조합, 구현, 이분탐색 정도일까?

여러 가지 알고리즘이 복합적으로 쓰이고, 완전탐색 -> 이분탐색 아이디어를 떠올리기도 쉽지 않은 문제였다.

개인적으로 Level 3 중에서 가장 어려운 편 아닌가 싶고, 3번 문제로 나온 건 난이도 조절 실패인 것 같다(개인적 의견).

백준으로 치면 골드 1 정도 될 거 같은 느낌.

 

from itertools import combinations, product

def binary_search(arr, target):
    """
    이분 탐색
    정렬된 점수 합 리스트를 받아, target보다 작은 요소의 개수를 반환합니다.
    """
    
    start = 0
    end = len(arr) - 1
    result = 0
    
    while start <= end:
        
        mid = (start + end) // 2
        
        if arr[mid] < target:
            result = mid + 1
            start = mid + 1
            
        else:
            end = mid - 1
    
    return result

def get_all_sums(dices):
    """
    주사위 리스트를 받아, 정렬된 모든 합 경우의 수 리스트를 리턴합니다.
    순서를 고려하는 product 입니다.
    """
    
    result = [sum(scores) for scores in product(*dices)]
    return sorted(result)

def solution(dice):
    """
    주사위 고르는 경우의 수 구하기 C(n,r) (combination)
    고른 주사위에 대한 점수 합 각각 구하기(6^r + 6^r) (product)
    이분 탐색으로 A가 이기는 횟수 구하기
    -> A의 모든 합을 순회하며, B합 리스트에 대해 이분 탐색
    """
    
    n = len(dice)
    
    # A가 선택한 모든 주사위 인덱스 리스트
    a_all_dices_idx = list(combinations(range(n), n//2))
    
    max_win_count = -1
    best_combination = None
    
    # A가 선택한 모든 주사위 경우의 수에 대해 반복
    for a_dices_idx in a_all_dices_idx:
        
        # B가 선택한 주사위 인덱스
        b_dices_idx = [i for i in range(n) if i not in a_dices_idx]
        
        # 인덱스를 실제 주사위 값으로 변환
        a_dices = [dice[i] for i in a_dices_idx]
        b_dices = [dice[i] for i in b_dices_idx]
        
        # 선택한 주사위로 모든 점수 합 경우의 수 구하기 (각 7776개)
        a_sums = get_all_sums(a_dices)
        b_sums = get_all_sums(b_dices)
        
        # 이분 탐색으로 A의 승리 횟수 구하기
        win_count = 0
        for a_sum in a_sums:
            win_count += binary_search(b_sums, a_sum)
        
        # 승리 횟수 최대값 갱신 및 리턴값 생성
        if win_count > max_win_count:
            max_win_count = win_count
            best_combination = a_dices_idx
    
    return sorted([i + 1 for i in best_combination])

 

A, B가 점수가 다른 주사위 n 개 중 각각 n / 2 개를 고르고, A가 승리할 확률이 가장 높은 주사위 조합을 고르는 문제이다.

n은 10 이하이고 주사위는 6개의 면을 가지고 있다. 만약 완전탐색으로 접근할 경우 무조건 시간초과가 난다고 보면 된다. 최대 연산횟수가 약 150억번 정도이기 때문이다.

 

이 문제는 크게 2개의 조합을 구해야 한다. 문제 해설은 주사위 10개(n = 10) 기준으로 하겠다.

n = 10

r = n // 2 (5)

1. 주사위 5개를 선택하는 경우의 수: C(n, r) -> 252개

2. A, B가 선택한 주사위에 대한 모든 점수 합 경우의 수: 6^r -> 7776개 (A, B 각각이므로 7776 * 2)

만약, 모든 경우의 수를 일일이 비교할 경우 시간 복잡도는 (C(n, r) * n^2) , 즉 N^2이 되고, 252 * (7776 * 7776) 이므로 약 150억번의 연산을 하게 된다.

 

따라서 A가 승리하는 횟수를 구하는 부분을 완전탐색이 아닌 이분탐색으로 개선해야 한다.

A, B의 점수 합 리스트를 따로 구한 후, A리스트를 순회하며(7776번) B리스트에서 몇 개의 요소가 A의 target보다 작은가를 세면 된다.

이러면 기존 완전탐색 N^2 -> 이분탐색 NlogN 으로 시간복잡도를 개선할 수 있게 된다.

 

모든 주사위 경우의 수(252)를 반복하되, 그 안에서 해당 주사위로 계산된 점수 합 리스트(7776)를 반복하며 이분탐색을 하는 것이다.

 

그리고 one based 형태로 주사위 인덱스를 1 더해 리턴해 주면 된다.