[이취코테] 큰 수의 법칙

문제

동빈이의 큰 수의 법칙은 다양한 수로 이루어진 배열이 있을 때 주어진 수들을 M번 더하여 가장 큰 수를 만드는 법칙이다.
단, 배열의 특정한 인덱스에 해당하는 수가 연속해서 K번 초과하여 더해질 수 없다. (서로 다른 인덱스에 해당하는 수가 같은 경우에도 서로 다른 것으로 간주한다.) 배열의 크기 N, 숫자가 더해지는 횟수 M, 그리고 K가 주어질 때 동빈이의 큰 수의 법칙에 따른 결과를 출력하시오. 

 

입력 조건:

• 첫째 줄에 N(2 <= N <= 1,000), M(1 <= M <= 10,000), K(1 <= K <= 10,000)의 자연수가 주어지며, 각 자연수는 공백으로 구분한다.
• 둘째 줄에 N개의 자연수가 주어진다. 각 자연수는 공백으로 구분한다. 단, 각각의 자연수는 1이상, 10,000 이하의 수로 주어진다.
• 입력으로 주어지는 K는 항상 M보다 작거나 같다.

 

출력 조건: 첫째 줄에 동빈이의 큰 수의 법칙에 따라 더해진 답을 출력한다.

 

나의 코드

n, m, k = map(int, input().split())
data = list(map(int, input().split()))
cnt = 0
result = 0

data.sort(reverse=True)
for i in range(m):
    if(cnt != k):
        result += data[0]
        cnt += 1
    else:
        result += data[1]
        cnt = 0
        
print(result)

이 문제는 전형적인 그리디 알고리즘 문제로 아이디어를 떠올리는 것은 어렵지 않았다.

풀이 방법은 가장 큰 수를 연속해서 K번 더하고 두 번째로 큰 수를 한 번 더하는 연산을 반복하는 것이다.

따라서 if ~ else 문을 사용하여 이를 구현하였다.

정답 코드

n, m, k = map(int, input().split())
data = list(map(int, input().split()))
data.sort()

first = data[n-1] # 가장 큰 수
second = data[n-2] # 두 번째 큰 수

count = int(m / (k+1)) * k
count += m % (k + 1)

result = 0
result += count * first
result += (m - count) * second

print(result)

정답 코드에서는 N이 100억 이상의 더 큰 수가 되는 경우를 가정했을 때도 시간초과가 나지 않도록 더 효율적으로 구현하였다.

가장 큰 수와 두 번째로 큰 수가 더해질 때 나타나는 특정한 수열을 이용하여 문제를 풀었다. 가장 큰 수가 K번 더해지고 두 번째로 큰 수가 1번 더해지는 것이 반복되므로 수열의 길이는 (K+1)이 된다. 따라서 M을 (K+1)로 나눈 몫이 수열이 반복되는 횟수가 되고, 다시 여기에 K를 곱해주면 가장 큰 수가 등장하는 횟수가 된다. 이 때 M이 (K+1)로 나누어떨어지지 않는 경우도 고려해줘야 하는데, 그럴 때는 M을 (K+1)로 나눈 나머지 만큼 가장 큰 수가 더해지므로 이를 고려하면 된다. 즉, 가장 큰 수가 더해지는 횟수는 다음과 같다.

count = int(m / (k+1)) * k + m % (k + 1)

 

 

배울 점

단순하게 구현할 수 있는 문제였지만 반복되는 수열을 생각해서 더 효율적으로 푸는 방법도 있다는 것을 알게 되었다.다음번에 비슷한 문제가 나오면 수열을 이용해서 풀 수 있도록 연습해야겠다.

 

---

참고자료: 이것이 취업을 위한 코딩 테스트다 with 파이썬 - 나동빈 지음(한빛미디어)