[백준][Python] 13251. 조약돌 꺼내기

문제

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

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

 

입력

첫째 줄에 M (1 ≤ M ≤ 50)이 주어진다.

둘째 줄에는 각 색상의 조약돌이 몇 개 있는지 주어진다. 각 색상의 조약돌 개수는 1보다 크거나 같고 50보다 작거나 같은 자연수이다.

셋째 줄에는 K가 주어진다. (1 ≤ K ≤ N)

 

출력

첫째 줄에 뽑은 조약돌이 모두 같은 색일 확률을 출력한다. 정답과의 절대/상대 오차는 10-9까지 허용한다.

---

코드

import sys

def multi(n, k):
    result = n
    for i in range(k-1):
       result = result * (n-1)
       n = n-1
    return result

if __name__ == '__main__':
    M = int(sys.stdin.readline().rstrip()) # 조약돌 색상 수
    
    temp = sys.stdin.readline().rstrip().split() # 각 색상 별 조약돌 개수
    list = []
    for i in temp:
        list.append(int(i))

    K = int(sys.stdin.readline().rstrip())  # 뽑을 조약돌 수
    
    N = sum(list)	# 전체 돌 수
    numerator = 0   # 분자
    denominator = multi(N, K) #분모
   
    # 만약 뽑을 값보다 적은 조약돌의 개수를 가진 색상은 필요없음
    # 해당 색상의 조약돌 제거
    for i in list:
        if(i < K):
            list.remove(i)
                  
    for i in list:
        numerator += multi(i, K)
    
    print(numerator/denominator)

 

풀이

문제의 이해를 위해 예제 입력 4로 설명을 한다.

예제 입력 4에서 조약돌의 색상(M)이 5가지이고, 각 색상의 조약돌의 개수는 [12,2,34,13,17]이며 조약돌을 랜덤하게 4개(K) 뽑는다고 되어있다 .

분모는 N개의 조약돌 중 K개를 뽑는 경우의 수이다. 따라서 C(N, K)를 구하면 된다. 이때 N은 조약돌의 전체 개수이다.

N은 12+2+34+13+17 = 78이 된다. 즉 분모는 C(78, 4)가 된다. 

분자는 각 색상별로 K를 뽑은 경우의 수를 더해주면 된다. 즉 C(한 색상의 조약돌 개수, K)를 모두 더해주면 된다.

따라서 C(12, 4)+C(2, 4)+C(34, 4)+C(13, 4)+C(17, 4)가 분자가 된다. 이때 C(2, 4)는 뽑을 조약돌 개수가 전체 조약돌 개수보다 적기 떄문에 불가능하다. 따라서 제외하면 C(12, 4)+C(34, 4)+C(13, 4)+C(17, 4) 가 분자가 된다.

 

∴ 확률값은 ${\frac{C(12, 4) + C(34,4) + C(13,4), C(17,4)}{C(78, 4)}}$ 가 된다.

 

계산을 해보자. 분자와 분모에 공통되는 수인 4!을 각각 분자와 분모에 곱해준다. 

분모의 경우는 78x77x76x75 = 34234200이 된다.

같은 방법으로 계산하면, 분자는 11880+1113029+17160+57120=1199184 가 된다.

 

따라서 분자/분모를 해주면 예시 출력 4와 같은 값이 나온다. 

 

이제 이걸 문제에 적용시키면 된다.

위에서 적었듯이, 분모는 C(N, K)가 되고 분자는 C(한 색상의 조약돌 개수, K)를 조약돌의 색상 개수별로 구해서 모두 더해주면 된다. 그리고 계산할 때는 K!을 곱했다고 생각하면 분모의 경우, K에 개수에 따라 Nx(N-1)x(N-2).. 이런식으로 구해주면 된다. 즉 K가 4면 Nx(N-1)x(N-2)x(N-3)을 해주면 되는 것이다.

 

코드에서 입력을 받은 후, N을 구해놓고 list를 돌아서 뽑을 조약돌의 개수보다 적은 조약돌의 개수를 가진 색상을 제거했다. 이는 분자에 포함되면 안되기 때문이다.

그리고 multi라는 함수에서 N과 K로 K!을 곱했다고 생각한 상황의 계산을 할 수 있게 했다. Combination 계산을 그냥 하고 싶다면 python의 라이브러리인 itertools에서 combinations을 import해 사용하면 된다.