[백준 20186] 수 고르기

수 고르기(20186)

문제 출처 : https://www.acmicpc.net/problem/20186

 

이 문제는 2020년 정보올림피아드 2차 대회 초등부 1번 문제 입니다.

문제 이해하기

문제를 읽어보면 상당히 복잡해 보입니다. 하지만 조금만 잘 생각해보면 규칙을 쉽게 찾을 수 있습니다.

2 3 1 2 1, K = 3

문제의 예제처럼 위와 같이 숫자가 있고 3개의 숫자를 선택해야 합니다. 그리고 점수는 자신의 왼쪽에 있는 선택된 수 입니다. 3개의 숫자를 어떻게 고를지 모르겠지만 K값이 3이기 때문에 숫자 3개를 골라야 합니다. 이 숫자들을 a, b, c라고 하겠습니다. a, b, c의 점수는 각각 a - 0, b - 1, c - 2 입니다. 이 숫자들의 합은 a + b + c - 0 - 1 - 2로 나타낼 수 있고 결국 a + b + c - 3이 됩니다. 즉 K값에 따라 전체 합계에서 빼야할 숫자는 정해져 있습니다. K 가 3이면 0 + 1 + 2로 3을 빼야 하고, K 가 5라면 0 + 1 + 2 + 3 + 4 로 10이 됩니다. K값에 따라 빼야 하는 숫자가 정해져 있기 때문에 골라야 하는 숫자 K개는 무조건 큰 숫자로 해야 합니다.

즉 이 문제는 K개의 가장 큰 숫자를 선택하고, 0부터 K - 1 까지의 합계를 빼면 되는 문제 입니다.

코드 작성하기

그럼 코드를 작성해 보겠습니다.

입력 받기

mii = lambda : map(int, input().split())
N, K = mii()
arr = list(mii())

골라야 할 자연수 전체 N개와, N개 중 골라야 할 자연수 K를 입력 받습니다. 그리고 자연수 N개를 arr이라는 리스트에 담았습니다.

계산하기

arr.sort()
ans = sum(arr[-K:]) - sum(range(K))
print(ans)

결과를 출력하기 위해서 알아야 할 값들은 다음과 같습니다. 숫자들 중 K개의 가장 큰 수, 0부터 K - 1까지의 합계 입니다. 정렬과 슬라이싱을 통해 K개의 가장 큰 수를 구했습니다. 그리고 range를 통해 0부터 K - 1개의 합계를 구했습니다.

0부터 K - 1까지의 합계는 등차수열의 합계로 쉽게 구할 수 있습니다. 초등부 문제이기 때문에 공식을 사용하기 보다는 range로 합계를 구해주었습니다.

마지막으로 K개의 가장 큰 수의 합계에서 0부터 K - 1까지의 합계를 빼주면 그 값이 이 문제에서 원하는 정답이 됩니다.

전체 코드

그럼 전체 코드를 알아 보겠습니다.

mii = lambda : map(int, input().split())
N, K = mii()
arr = list(mii())

arr.sort()
ans = sum(arr[-K:]) - sum(range(K))
print(ans)