주사위 고르기 (프로그래머스, 2024 카카오 겨울 인턴십, Python)

풀이

(1) 아이디어

완전탐색

주사위 10개에 대해 완전탐색을 한다면, 먼저 주사위 10개를 5개씩 나눠갖고 (10 C 5), A가 먼저 5개의 주사위를 굴리고 (6**5), 그 다음 B가 5개의 주사위를 굴리고 (6**5), 각각의 모든 경우에 대해 주사위의 합을 비교하는 경우의 수(10 C 5 * 6**10)를 계산해야 합니다. 하지만 이 경우 시간초과가 발생합니다.

이분탐색

주사위를 절반으로 나눠갖고, A와 B가 각각 주사위를 굴리는 것까지는 동일하지만 주사위의 합을 비교하는 과정에서 이분탐색으로 시간을 줄일 수 있습니다. 먼저 각자 고른 주사위로 만들 수 있는 주사위의 합을 각자 리스트에 저장합니다. 그리고 B 리스트를 정렬합니다. 이제 A 리스트 각각의 원소를 이용해서 B 리스트에서 leftmost 이분탐색을 실시합니다. 이분탐색으로 찾은 인덱스보다 왼쪽에 있는 B의 원소들은 현재 고른 A의 원소보다 작습니다. 따라서 여기서 찾은 인덱스가 A가 이길 수 있는 경우의 수가 됩니다.

#				5						# A 리스트의 원소
[0, 1, 2, 3, 4, 5, 5, 5, 5, 6, 7, 8, 9] # 정렬된 B 리스트

예를 들어 A 리스트의 원소를 5라고 해보겠습니다. B 리스트에서 5에 대해 leftmost 이분탐색을 실시합니다. 이렇게 찾아진 5의 인덱스는 왼쪽에 있는 B의 원소들보다 무조건 큽니다. 따라서 5로 이길 수 있는 경우의 수는 5가지입니다.

#							   7		# A 리스트의 원소
[0, 1, 2, 3, 4, 5, 5, 5, 5, 6, 7, 8, 9] # 정렬된 B 리스트

이번에는 A 리스트의 원소를 7이라고 해보겠습니다. B 리스트에서 7에 대해 leftmost 이분탐색을 실시합니다. 이렇게 찾아진 7의 인덱스는 왼쪽에 있는 B의 원소들보다 무조건 큽니다. 따라서 7로 이길 수 있는 경우의 수는 10가지입니다.

(2) 변수 정의

    dic = {}
    L = len(dices)

dic의 key는 이길 수 있는 경우의 수, value는 그때의 주사위 번호들입니다. L은 주사위의 개수입니다.

(3) 주사위 고르기

    for A_index_combi in combinations(range(L), L//2):
        B_index_combi = [i for i in range(L) if i not in A_index_combi]

먼저 A가L개의 주사위 중에서 L//2개의 주사위를 가져갑니다. B는 A가 고르지 않은 나머지 주사위를 가져갑니다.

(4) 만들 수 있는 주사위의 합 구하기

        A, B = [], []
        for order_product in product(range(6), repeat=L//2):
            A.append(sum(dices[i][j] for i, j in zip(A_index_combi, order_product)))
            B.append(sum(dices[i][j] for i, j in zip(B_index_combi, order_product)))
        B.sort()

주사위 L//2개를 각각 굴려서 만들어지는 주사위 면들의 조합 order_product를 각자가 고른 주사위 번호들의 조합과 패킹하고, 이를 이용하여 각자의 주사위 조합으로 만들 수 있는 모든 주사위의 합을 구합니다. B 리스트의 경우 이분탐색을 위해 정렬을 수행합니다.

(5) 이분탐색

        wins = sum(bisect_left(B, num) for num in A)
        dic[wins] = list(A_index_combi)

A의 각 원소들로 B 리스트에서 leftmost 이분탐색을 실시합니다. 여기서 반환되는 인덱스 값이 해당 A의 원소로 B 리스트에서 이길 수 있는 경우의 수입니다. 그리고 이렇게 찾아진 모든 경우의 수를 더한 값이 현재 주사위 조합으로 이길 수 있는 모든 경우의 수 wins입니다. wins를 key로 하여 dic 딕셔너리에 A 주사위 조합을 저장합니다. 이때 같은 key를 갖는 다른 조합이 있을 수도 있지만, 문제 조건에 의해 가장 승률이 높은 주사위 조합은 유일하므로 무시해도 괜찮습니다.

(6) 결과 반환

	max_key = max(dic.keys())

    return [x+1 for x in dic[max_key]]

승수의 최댓값 max_key를 찾고, 해당 key에 대한 주사위 번호들의 리스트를 반환합니다. 단, 여기에 저장된 주사위 번호는 0부터 시작하므로, 1씩 더해서 반환합니다.

소스코드

from itertools import combinations, product
from bisect import bisect_left

def solution(dices):
    dic = {}
    L = len(dices)
    for A_index_combi in combinations(range(L), L//2):
        B_index_combi = [i for i in range(L) if i not in A_index_combi]
    
        A, B = [], []
        for order_product in product(range(6), repeat=L//2):
            A.append(sum(dices[i][j] for i, j in zip(A_index_combi, order_product)))
            B.append(sum(dices[i][j] for i, j in zip(B_index_combi, order_product)))
        B.sort()

        wins = sum(bisect_left(B, num) for num in A)
        dic[wins] = list(A_index_combi)

    max_key = max(dic.keys())

    return [x+1 for x in dic[max_key]]