[이코테/그리디] 볼링공 고르기 (정답, 풀이)

문제

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 = map(int,input().split())
ball = list(map(int,input().split())) 

cnt = 0
for i in range(len(ball)):
    for j in range((len(ball))):
        if i == j or ball[i] == ball[j]:
            break
        cnt+=1

print(cnt//2)

 

나는 이중 for문을 사용해 같은 인덱스(같은 볼링공)가 아니고, 무게가 같지 않다면 카운트를 증가시켰다.

그리고 마지막에 (1,2) , (2,1) 같은 순서만 다르고 같은 조합을 제외하기 위해 cnt를 2로 나누어서 답을 구했다.

 

틀린 답은 아니라고 생각하지만 문제에서 제시한 볼링공의 최대 개수가 1000개이기 때문에

내 코드는 최악의 경우 1,000,000 번의 연산을 하게 된다. 

 

너무 쉽게 풀었지만 기본적인 시간복잡도를 놓쳐버림..

 

정답

n,m = map(int,input().split())
ball = list(map(int,input().split()))
cnt = 0

# 각 인덱스의 무게를 가진 볼링공이 몇 개 있는지 저장
wgt = [0] * 11
for i in ball:
    wgt[i] += 1

for i in range(1,m+1):
    n -= wgt[i]  
    cnt += n * wgt[i]
    # wgt[i] -> a가 선택할 공
    # n -> b가 선택할 공
 
print(cnt)

 

 

풀이

답지를 보고 원리는 쉽게 이해를 했는데 소스코드를 이해하기가 조금 어려웠다. 

 

원리는 이렇다. 

무게가 각각 1, 3, 2, 3, 2  인 공이 5개 있다면

먼저 무게별로 공의 개수를 센다. 

 

무게  -  개수

  1     -   1개

  2     -   2개

  3     -   2개

 

이 때, a와 b가 공을 차례로 선택하는 데 

먼저 a가 무게가 1인 공을 선택할 때 b가 나머지 공을 선택하는 경우의 수는

(무게가 1인 공의 개수 : 1개) * (무게가 2인 공의 개수 2개 + 무게가 3인 공의 개수 2개) =

1  *  4  =  4 이다.

 

그 다음, a가 무게가 2인 공을 선택할 때 b가 나머지 공을 선택하는 경우의 수는

(무게가 2인 공의 개수 : 2개) * (무게가 3인 공의 개수 : 2개) = 

2  *  2  =  4 이다.

 

마지막으로, a가 무게가 3인 공을 선택할 때 b가 나머지 공을 선택하는 경우의 수는 

(무게가 3인 공의 개수: 2개) * (없음)

3  *  0  =  0

 

결과적으로, 4 + 4  = 8 가지 경우의 수가 존재함을 알 수 있다. 

여기서 포인트는 a가 공을 선택하고 나서 b가 선택할 수 있는 공의 선택지는 줄어든다는 것이다. 

 

---

 

이 원리를 코드에 적용시켜 보면 다음과 같다.

 

wgt라는 리스트를 만들어서 공의 무게별 개수를 담는다.

위의 예제를 참고하면

wgt 라는 리스트에는 데이터가 [ 0, 1, 2, 2, 0, 0 ] 이렇게 들어가 있을 것이다

(1번 인덱스에 1의 무게를 가진 공 1개, 2번 인덱스에 2의 무게를 가진 공 2개, 3번 인덱스에 3의 무게를 가진 공 2개...)

 

그럼 for문을 돌면서 제시된 무게 m+1까지 (m번째 인덱스를 사용해야하니 m+1까지)

a가 선택할 수 있는 경우의 수와 b가 선택할 수 있는 경우의 수를 곱하여 결과값에 더해주면 끝

 

for i in range(1,m+1):
    n -= wgt[i] 
    cnt += n * wgt[i]
    # wgt[i] -> a가 선택할 공
    # n -> b가 선택할 공

공의 무게를 i로 잡고 for문을 돌며 

i번 공을 선택했을 때 a가 선택할 수 있는 경우의 수인 (=i번 공의 개수) wgt[i] 와

그 후 b가 선택할 수 있는 경우의 수인(전체 공 개수- i번 공의 개수) n-=wgt[i]를 곱하여 cnt에 계속 더한다.

 

 

 

언제쯤 그리디를 답지없이 완벽하게 풀 수 있을까 ㅠ

 

 

참조

이것이 취업을 위한 코딩테스트다.(with. python)