< SW 마에스트로 코딩테스트-python > 볼링공 고르기

풀이 시간 : 30분

시간 제한 : 1초

메모리 제한 : 128MB

출처 : 2019 SW 마에스트로 입학 테스트 (이것이 취업을 위한 코딩테스트다 - 나동빈)

문제

A, B 두 사람이 볼링을 치고 있습니다. 두 사람은 서로 무게가 다른 볼링공을 고르려고 합니다. 볼링공은 총 N 개가 있으며 각 볼링공마다 무게가 적혀 있고, 공의 번호는 1번부터 순서대로 부여됩니다. 또한 같은 무게의 공이 여러 개 있을 수 있지만, 서로 다른 공으로 간주합니다. 볼링공의 무게는 1부터 M까지의 자연수 형태로 존재합니다.

예를 들어 N이 5이고, M이 3이며 각각의 무게가 차례대로 1, 3, 2, 3, 2일 때 각 공의 번호가 차레대로 1번부터 5번까지 부여됩니다. 이때 두 사람이 고를 수 있는 볼링공 번호의 조합을 구하면 다음과 같습니다.

(1번, 2번), (1번, 3번), (1번, 4번), (1번, 5번), (2번, 3번), (2번, 5번), (3번,
4번), (4번, 5번)

결과적으로 두 사람이 공을 고르는 경우의 수는 8가지입니다. N개의 공의 무게가 각각 주어질 때, 두 사람이 볼링공을 고르는 경우의 수를 구하는 프로그램을 작성하세요.

입력 조건

• 첫째 줄에 볼링공의 개수 N, 공의 최대 무게 M이 공백으로 구분되어 각각 자연수 형태로 주어집니다. (1 ≤ N ≤ 1,000, 1 ≤ M ≤ 10)
• 둘째 줄에 각 볼링공의 무게 K가 공백으로 구분되어 순서대로 자연수 형태로 주어집니다. (1 ≤ K ≤ M)

출력 조건

• 첫째 줄에 두 사람이 볼링공을 고르는 경우의 수를 출렵합니다.

내가 적어 보는 문제 풀이

문제를 보자마자 떠올린 방법은 조합이다. 순서를 무시하기 때문에 순열은 아니지만, 각 볼링공에 번호가 있고 그 번호가 볼링공의 고유 이름처럼 쓰이고 있기 때문에 서로 다른 N개 중에 2개를 뽑는 경우의 수가 떠올랐다.

단, 볼링공의 무게는 중복된 무게가 존재할 수 있는데, 중복되는 무게는 뽑지 않는다고 했다. 예를 들어 무게가 3인 볼링공을 두 사람이 같이 뽑지는 않기 때무에 이 경우들을 제외해 줘야 한다.

구해야 하는 답이 볼링공을 뽑은 경우들을 나열하는 것이 아니라 경우의 수만 구하면 되기 때문에 일일히 반복문을 돌기 보다는 할 수 있다면 수학적 공식으로 수만 빠르게 구하는 것이 훨씬 효울적이다.

nCr = n!/(n-r)!r!

위 조합 공식을 이용해 먼저 볼링공 2개를 뽑는 전체 경우의 수를 구한다.

from math import factorial

n, m  = map(int, input().rstrip().split())
li = list(map(int, input().rstrip().split()))
result = factorial(n) // (factorial(n-2)*2)  # 주의: ()로 묶어야 함

이제 result에서 같은 무게를 뽑은 경우의 수를 빼주면 된다.

그러기 위해 li에서 무게들을 카운팅 해, 개수가 2이상인 무게들이 각각 이룰 수 있는 조합의 수를 구해야 한다.

가장 편리한 방법은 collections의 Counter()를 쓰면 된다.

https://docs.python.org/ko/3/library/collections.html#collections.Counter

from collections import Counter

# ...

count = Counter(li)
for i in range(1, m+1):
    n = count[i]
    if n == 1 or n == 0:  # 1개나 0개는 예외 처리
        continue
    if n == 2:  # n == 2 이면 분모가 0이 되기 때문에 따로 경우의 수 1을 빼줌
        result -= 1
        continue
    result -= factorial(n) // factorial(n-2)*2

입출력 결과

>>> 5 3
>>> 1 3 2 3 2
8

>>> 8 5
>>> 1 5 4 3 2 4 5 2
25

전체 코드 - https://github.com/on1ystar/algorithm-problem-solving-codes/blob/master/greedy/볼링공 고르기.py