주사위 고르기

from itertools import combinations, product def solution(dice): ''' A가 이기는 것을 기준으로 함. ''' answer = [] dice_indices = [i for i in range(len(dice))] # 주사위 인덱스 pickup_num = int(len(dice) / 2) # 뽑을 주사위 갯수 win_max = 0 for a_dices in combinations(dice_indices, pickup_num): a_dices = set(a_dices) # A가 뽑은 주사위 집합 (ex: (1, 2)) b_dices = set(dice_indices) - a_dices # B가 뽑은 주사위 집합 (ex: (3, 4)) ### A와 B 따로따로 계산 수행 ### # A, B 각각 나올 수 있는 주사위 눈의 조합 itemsA = [dice[a_idx] for a_idx in a_dices] # ex 2 itemsB = [dice[b_idx] for b_idx in b_dices] # ex 2 product_listA = list(product(*itemsA)) # 6^2 product_listB = list(product(*itemsB)) # 6^2 winA = 0 # A가 승리할 수 있는 방법의 수 # A, B 각각 주사위 눈의 조합으로 얻을 수 있는 숫자 합 sum_tempA, sum_tempB = [], [] # A에 대하여, 가능한 주사위 눈의 숫자 합을 저장 for a in product_listA: sum_tempA.append(sum(list(a))) # B에 대하여, 가능한 주사위 눈의 숫자 합을 저장 for b in product_listB: sum_tempB.append(sum(list(b))) # A, B에 대한 주사위 숫자 합 정렬 sum_tempA = sorted(sum_tempA) sum_tempB = sorted(sum_tempB) for target in sum_tempA: # target이 B의 주사위 눈 합의 최솟값보다 작은 경우 if target <= sum_tempB[0]: continue # target이 B의 주사위 눈 합의 최댓값보다 큰 경우 if target > sum_tempB[-1]: winA += len(sum_tempB) continue ### 이분 탐색 수행 ### left, right = 0, len(sum_tempB) - 1 find = 0 while left <= right: mid = (left + right) // 2 if target > sum_tempB[mid]: left = mid + 1 elif target < sum_tempB[mid]: right = mid - 1 # target을 B에서 찾아도, 동일한 값이 이전 인덱스에도 존재할 수 있습니다. # 우리는 A가 이기는 경우(점수가 높은 경우)만을 고려해야 합니다. # 따라서, right를 줄여서 계속 이분 탐색을 수행해야 함을 잊지 마세요! else: find = mid right = mid - 1 # A가 이길 수 있는 방법의 수 갱신 if find != 0: winA += find else: winA += (right + 1) # 갱신된 방법의 수가 최댓값인 경우 if winA > win_max: answer = [idx+1 for idx in a_dices] win_max = winA return sorted(answer)