주사위 고르기
2024 카카오 윈터 인턴쉽 문제가 오픈되어 풀어보았다.
문제 설명
- n 개의 주사위가 있다.
- 1 ~ n 까지의 숫자가 적혀있다.
- 주사위 구성이 주어지는데 이 때 어떤 주사위 구성을 선택해야 이길 확률이 가장 높은지 구하라
이길 확률이 가장 높은 주사위 구성을 반환해라
흐름 파악
흐름 파악을 위해 꼬리에 꼬리를 무는 방식으로 접근하려한다.
이길 확률이 가장 높은 주사위 구성은 어떻게 구할 수 있을까?
우선 주사위 구성은 전체 구성에서 A 와 B 가 절반씩 나눠가지게돼.
즉 A 가 가져간 주사위 구성을 제외하면 B 가 가져간 주사위 구성이 돼.
이긴다는건 뭐야?
sum(A i..j) > sum(B i..j) 이면 A 가 이기는거지
이길 확률은 어떻게 구할 수 있을까?
모든 경우의 수를 구해서 이길 확률을 구하면 되겠네
로직은 간단한데 성능이 문제가 될 것 같은데?
주사위 구성은 최대 10개이고 주사위 구성의 길이는 6개이면서 원소의 길이는 100개야.
그러면 10C5 * 6^6 이야.
252 * 46656 = 11,757,312 이야.
우선 이렇게 구현해보고 성능을 봐야겠어.
1차 구현
from itertools import combinations, product
def solution(dice):
answer = []
max_possible = [[], 0]
pick_combinations = list(combinations([i for i in range(len(dice))], len(dice) // 2))
i = 0
while i < len(pick_combinations):
a_dice = list(pick_combinations[i])
b_dice = list(set(i for i in range(len(dice))) - set(list(a_dice)))
a_comb = sorted(list(map(sum, product(*[dice[d] for d in a_dice]))))
b_comb = sorted(list(map(sum, product(*[dice[d] for d in b_dice]))))
win_count = 0
for a in a_comb:
for b in b_comb:
if a > b:
win_count += 1
continue
break
if win_count > max_possible[1]:
max_possible[0] = list(map(lambda x: x + 1, a_dice))
max_possible[1] = win_count
i += 1
return max_possible[0]
결과
테스트케이스가 18번까지만 정답이고 19번부터는 시간초과가 발생한다. 그래도 논리가 틀리지는 않았다.
2차 구현
1차 구현 코드에서 A[] 와 B[] 가 있을 때 두 배열 모두 정렬이 되어있다. 이 경우 A 의 원소가 B[-1] 원소보다 크면 A 의 전승이다. 이것을 토대로 로직을 수정해보았다.
if a > b_comb[-1]:
win_count += len(b_comb)
continue
결과
비교해보면 아주 조금 빨라졌기는하지만 여전히 18번 테스트케이스까지만 정답이다.
3차 구현
예를 들어 배열 A 는 [2, 2, 3, 3, 4, 4, 4, 4, 5, 6, 6, 11, 12, 12, 12] 처럼 중복된 숫자가 있다. 또한 같은 숫자인 A[0] 이나 A[1] 이나 같은 승수를 가져가게 될 것이다. 그러므로 숫자에 대한 승수 기록을 가지고 있으면 for 을 돌지 않아도 되겠다.
win_by_number_dict 변수를 추가해서 기록이 있으면 for 돌지 않고 바로 승수를 추가하고 기록이 없으면 계산하도록 하였다.
win_by_number_dict = defaultdict(int)
for a in a_comb:
count = 0
if a in win_by_number_dict:
win_count += win_by_number_dict[a]
continue
...
win_by_number_dict[a] = count
결과
드디어 22, 23, 24 를 제외한 테스트케이스가 통과하였다. 그러나 여전히 시간 초과가 발생한다.
4차 구현
B 가 [3, 3, 4, 4, 4, 11, 11, 12] 라는 값이 있을 때 a 가 B 의 4 를 이기면 4를 포함해 3까지 이기게 된다. 또 11을 이기게 되면 11을 포함해 11이하의 수까지 이기게 된다. 이런식으로 계산하면 시간을 줄일 수 있을 것 같다.
즉 B 의 원소마다 자신을 포함하고 자신보다 작은 숫자의 개수를 저장한다.
B 는 정렬을 내림차순으로 변경해야하는데 이유는 가장 큰 수부터 비교하다가 이겼을 때 누적된 승수를 더할 수 있기 때문이다.
b_comb = sorted(list(map(sum, product(*[dice[d] for d in b_dice]))), reverse=True)
acc = 0
acc_dict = defaultdict(int)
for k, v in b_counter:
acc += v
acc_dict[k] = acc
win_count = 0
...
for b in b_comb:
if a > b:
count += acc_dict[b]
break
win_by_number_dict[a] = count
win_count += count
결과
23번 케이스를 제외하고는 성공을 한다. 다른 통과 시간을 보니 타임아웃은 10초인 것 같다.
5차 구현 - 성공
생각해보면 A 배열과 B 배열의 이중 반복문이 문제임이 확실하다. 그래서 어떻게 하면 반복의 수를 줄일 수 있을까 생각하다 B 의 누적된 승수를 저장하고 있으니 set(B) 로 사용해도 되겠다는 생각이 들었다.
그렇게 되면 중복된 숫자가 제거되니 B 배열의 원소 개수는 굉장히 줄어들 것이다.
b_comb = list(map(sum, product(*[dice[d] for d in b_dice])))
b_counter = sorted(list(Counter(b_comb).items()))
b_comb = sorted(list(set(b_comb)), reverse=True)
결과 확인
모든 테스트 케이스가 통과하였다.
이전 구현과 22번을 비교해보면 9049ms -> 1281ms 로 7배 이상 빨라졌다.
최종 코드
from collections import defaultdict, Counter
from itertools import combinations, product
def solution(dice):
max_possible = [[], 0]
pick_combinations = list(combinations([i for i in range(len(dice))], len(dice) // 2))
i = 0
while i < len(pick_combinations):
a_dice = list(pick_combinations[i])
b_dice = list(set(i for i in range(len(dice))) - set(list(a_dice)))
a_comb = sorted(list(map(sum, product(*[dice[d] for d in a_dice]))))
b_comb = list(map(sum, product(*[dice[d] for d in b_dice])))
b_counter = sorted(list(Counter(b_comb).items()))
b_comb = sorted(list(set(b_comb)), reverse=True)
acc = 0
acc_dict = defaultdict(int)
for k, v in b_counter:
acc += v
acc_dict[k] = acc
win_count = 0
win_by_number_dict = defaultdict(int)
for a in a_comb:
count = 0
if not a > b_comb[-1]:
win_by_number_dict[a] = 0
win_count += 0
continue
if a > b_comb[0]:
win_by_number_dict[a] = len(a_comb)
win_count += len(a_comb)
continue
if a in win_by_number_dict:
win_count += win_by_number_dict[a]
continue
for b in b_comb:
if a > b:
count += acc_dict[b]
break
win_by_number_dict[a] = count
win_count += count
if win_count > max_possible[1]:
max_possible[0] = list(map(lambda x: x + 1, a_dice))
max_possible[1] = win_count
i += 1
return max_possible[0]