백준 20186 수 고르기 Python: K개 최댓값과 등차수열 벌점

백준 20186 수 고르기의 답은 가장 큰 수 K개의 합에서 K(K-1)/2를 뺀 값이다. 어떤 위치의 수를 고르더라도 선택된 수들이 받는 전체 벌점은 항상 0 + 1 + ... + (K-1)이기 때문이다.

따라서 위치 조합을 탐색하거나 DP를 만들 필요가 없다. 수열을 정렬해 가장 큰 K개를 고르면 된다.

벌점은 왜 선택 위치와 무관할까

선택한 K개의 위치를 왼쪽부터 정렬해 보자. 가장 왼쪽에 있는 선택된 수의 왼쪽에는 선택된 수가 0개다. 두 번째 수의 왼쪽에는 1개, 세 번째 수의 왼쪽에는 2개가 있다.

마지막 선택 수까지 쓰면 벌점은 다음과 같다.

0, 1, 2, ..., K-1

합은 등차수열 공식으로 정해진다.

0 + 1 + ... + (K-1) = K(K-1) / 2

이 값은 어떤 위치를 선택해도 같다. 전체 점수를 키우려면 벌점이 아니라 선택한 원래 수의 합만 최대화하면 되므로, 가장 큰 수 K개가 최적이다.

예제 2 3 1 2 1에서 K=3이면 큰 수 세 개는 3, 2, 2다.

선택한 수의 합 = 7
전체 벌점 = 0 + 1 + 2 = 3
최대 점수 = 7 - 3 = 4

Python 풀이

n, k = map(int, input().split())
numbers = list(map(int, input().split()))

numbers.sort(reverse=True)

penalty = k * (k - 1) // 2
answer = sum(numbers[:k]) - penalty

print(answer)

나눗셈에는 /가 아니라 //를 사용했다. K(K-1)은 연속한 두 정수의 곱이라 항상 짝수이므로 결과는 정확한 정수다. int(k * (k - 1) / 2)처럼 실수로 바꿨다가 다시 정수로 변환할 이유가 없다.

문제 제한은 1 ≤ N ≤ 5,000, 1 ≤ K ≤ N, 각 수는 1 이상 100,000 이하다. 선택한 수의 합은 최대 5억이고 벌점도 약 1,250만 이하라 이 문제의 범위에서는 32비트 부호 있는 정수로도 계산 가능하다. Python 정수는 이보다 큰 값도 정확하게 다룬다.

시간 복잡도와 heap 선택 기준

정렬이 대부분의 시간을 차지하므로 시간 복잡도는 O(N log N)이다. 입력 배열과 정렬 결과를 저장하는 공간도 필요하다.

heapq.nlargest(k, numbers)를 쓰면 일반적으로 O(N log K)에 큰 수 K개를 구할 수 있다. 다만 이 문제는 N이 최대 5,000이라 전체 정렬 풀이가 충분히 단순하고 빠르다. K가 N에 가깝다면 heap의 이점도 작다. 복잡도 표기만 보고 무조건 heap이 낫다고 결론 내리기보다 입력 크기와 구현 단순성을 함께 보는 편이 좋다.

풀이를 어떻게 검증했나

작은 N에서는 가능한 K개 위치 조합을 모두 열거할 수 있다. 각 조합의 점수를 직접 계산한 값과 정렬 공식을 비교했다.

from itertools import combinations


def brute_force(numbers, k):
    best = -10**18

    for picked in combinations(range(len(numbers)), k):
        score = sum(
            numbers[index] - rank
            for rank, index in enumerate(picked)
        )
        best = max(best, score)

    return best

N=1부터 8까지 무작위 자연수로 만든 3,600개 작은 입력에서 완전탐색 결과와 sum(sorted(numbers, reverse=True)[:k]) - k*(k-1)//2가 모두 일치했다.

참고 자료

• 백준 20186 수 고르기