BOJ 13251 - 조약돌 꺼내기 (Python)

S3, 조합

문제

효빈이의 비밀 박스에는 조약돌이 N개 들어있다. 조약돌의 색상은 1부터 M까지 중의 하나이다.

비밀 박스에서 조약돌을 랜덤하게 K개 뽑았을 때, 뽑은 조약돌이 모두 같은 색일 확률을 구하는 프로그램을 작성하시오.

풀이

• 처음에 combinations()를 통해 풀었는데 시간초과가 났다.

from sys import stdin
from itertools import combinations

m = int(stdin.readline())
arr = list(map(int, stdin.readline().split()))
k = int(stdin.readline())

stones = []
i = 1
for num in arr:
    stones.extend([i]*num)
    i += 1

total = 0
cnt = 0

for tmp in combinations(stones, k):
    if len(set(tmp)) == 1:
        cnt += 1
    total += 1
print(cnt/total)

• 해당 돌만 뽑는 경우의 수들의 합 / 전체 돌 중에 k개 뽑는 경우로 해결할 수 있다.

from sys import stdin
import math

m = int(stdin.readline())
arr = list(map(int, stdin.readline().split()))
k = int(stdin.readline())

n = sum(arr)
total = math.comb(n, k)

color = 0
for stone in arr:
    color += math.comb(stone, k)

print(color/total)

수학으로 풀이를 생각했어야 됬다.