문제 요약
n개의 주사위가 주어진다.
A가 먼저 n / 2개의 주사위를 고르고, B는 남은 n / 2개의 주사위를 가진다.
각자 가진 주사위를 모두 굴려 나온 수의 합으로 승부한다.
목표는 A가 승리할 확률이 가장 높아지는 주사위 조합을 찾는 것이다.
정답은 A가 골라야 하는 주사위 번호를 오름차순으로 반환한다.
핵심 아이디어
A가 고를 수 있는 주사위 조합을 모두 확인한다.
어떤 조합을 A가 선택하면, B의 조합은 자동으로 결정된다.
각 조합에 대해 다음을 계산한다.
- A가 만들 수 있는 모든 합
- B가 만들 수 있는 모든 합
- A의 합이 B의 합보다 큰 경우의 수
가장 승리 횟수가 큰 A의 조합이 정답이다.
모든 합 구하기
주사위 여러 개를 굴렸을 때 가능한 합을 구해야 한다.
처음에는 합이 0인 경우 하나에서 시작한다.
이후 주사위를 하나씩 추가하면서 가능한 합을 갱신한다.
sums = [0]
for index in selected:
next_sums = []
for current in sums:
for value in dice[index]:
next_sums.append(current + value)
sums = next_sums이렇게 하면 선택한 주사위들로 만들 수 있는 모든 합을 구할 수 있다.
승리 횟수 세기
A의 합 목록과 B의 합 목록을 모두 만든 뒤, A가 이기는 경우를 세어야 한다.
단순히 모든 쌍을 비교하면 시간이 오래 걸릴 수 있다.
따라서 B의 합 목록을 정렬한 뒤, A의 각 합보다 작은 B의 합 개수를 이분 탐색으로 구한다.
bisect_left(b_sums, a_sum)bisect_left는 정렬된 배열에서 a_sum이 들어갈 가장 왼쪽 위치를 반환한다.
즉, a_sum보다 작은 값의 개수를 의미한다.
이 값이 곧 해당 a_sum으로 A가 이기는 경우의 수다.
풀이 과정
1. A가 고를 주사위 조합 생성
combinations(range(n), n // 2)주사위 index 기준으로 n / 2개를 고르는 모든 조합을 만든다.
2. B의 주사위 조합 구하기
A가 고른 주사위를 제외한 나머지가 B의 주사위다.
b_indices = [i for i in range(n) if i not in a_set]3. A와 B의 합 분포 생성
각 조합에 대해 가능한 모든 합을 만든다.
a_sums = get_sums(a_indices)
b_sums = get_sums(b_indices)4. B의 합 정렬 후 이분 탐색
b_sums.sort()이후 A의 각 합마다 이기는 B의 합 개수를 누적한다.
5. 최대 승리 횟수 갱신
승리 횟수가 가장 큰 조합을 정답으로 저장한다.
문제에서 최적 조합은 유일하다고 했으므로 동점 처리를 따로 고민하지 않아도 된다.
Python 코드
from itertools import combinations
from bisect import bisect_left
def solution(dice):
n = len(dice)
half = n // 2
best_win_count = -1
best_combination = None
def get_sums(indices):
sums = [0]
for index in indices:
next_sums = []
for current in sums:
for value in dice[index]:
next_sums.append(current + value)
sums = next_sums
return sums
for a_indices in combinations(range(n), half):
a_set = set(a_indices)
b_indices = [i for i in range(n) if i not in a_set]
a_sums = get_sums(a_indices)
b_sums = get_sums(b_indices)
b_sums.sort()
win_count = 0
for a_sum in a_sums:
win_count += bisect_left(b_sums, a_sum)
if win_count > best_win_count:
best_win_count = win_count
best_combination = a_indices
return [index + 1 for index in best_combination]코드 설명
조합 생성
for a_indices in combinations(range(n), half):A가 선택할 수 있는 모든 주사위 조합을 확인한다.
주사위 번호는 문제에서 1번부터 시작하지만, 코드에서는 0-index로 처리한다.
마지막에 정답을 반환할 때 1을 더한다.
가능한 합 계산
def get_sums(indices):선택한 주사위들로 만들 수 있는 모든 합을 반환한다.
주사위를 하나씩 추가하면서 기존 합에 현재 주사위의 각 면 값을 더한다.
B의 조합 계산
a_set = set(a_indices)
b_indices = [i for i in range(n) if i not in a_set]A가 선택한 주사위를 제외한 나머지 주사위가 B의 주사위다.
승리 횟수 계산
win_count += bisect_left(b_sums, a_sum)정렬된 b_sums에서 a_sum보다 작은 값의 개수를 구한다.
동점은 승리가 아니므로 a_sum보다 작은 경우만 세어야 한다.
따라서 bisect_left를 사용한다.
정답 갱신
if win_count > best_win_count:승리 횟수가 더 큰 조합을 찾으면 정답을 갱신한다.
승리 확률은 전체 경우의 수가 모든 조합에서 같으므로, 승리 횟수를 비교하면 된다.
시간 복잡도
주사위 개수를 n이라고 하자.
A가 고르는 조합 수는 다음과 같다.
C(n, n / 2)각 조합마다 가능한 합의 개수는 다음과 같다.
6^(n / 2)B의 합을 정렬하고, A의 각 합에 대해 이분 탐색을 수행한다.
따라서 전체 시간 복잡도는 대략 다음과 같다.
O(C(n, n / 2) * 6^(n / 2) * log(6^(n / 2)))제한에서 n이 크지 않기 때문에 가능한 방식이다.
공간 복잡도
A와 B의 합 목록을 저장한다.
O(6^(n / 2))정리
이 문제는 주사위 조합을 완전 탐색하되, 승리 횟수 계산을 효율적으로 해야 한다.
핵심은 다음과 같다.
- A가 고를 수 있는
n / 2개 조합을 모두 확인한다. - 각 조합에 대해 A와 B가 만들 수 있는 모든 합을 구한다.
- B의 합을 정렬한다.
- A의 각 합마다 이분 탐색으로 이기는 경우의 수를 센다.
완전 탐색과 이분 탐색을 함께 사용하면 깔끔하게 해결할 수 있다.
