문제 88

균형 잡힌 수열 만들기

내 상태: 미제출
난이도: 88번 문제 난이도 보기

Platinum IV
출제자: rlatjwls7882
시간 제한: 1000 ms
메모리 제한: 512 MB

태그 보기

그리디 알고리즘정렬

양의 정수로 이루어진 길이 $N$의 수열 $A$가 주어진다.

수열의 불균형도를 수열의 원소의 최댓값과 최솟값의 차로 정의한다.

세종이는 다음 연산을 한 번 이상 수행해야 한다.

수열 $A$에서 $1$개 이상 $K$개 이하의 원소를 선택한다.
선택하지 않은 원소들의 합을 $S$라고 할 때, 선택한 각 원소에 $S$를 더한다.

세종이가 만들 수 있는 수열의 불균형도의 최솟값을 구하라.

입력

첫 번째 줄에 두 정수 $N$, $K$가 공백으로 구분되어 주어진다. $(2 \le N \le 200\,000; 1 \le K<N)$

두 번째 줄에 $N$개의 정수 $A_1, A_2, \cdots, A_N$이 공백으로 구분되어 주어진다. $(1 \le A_i \le 10^9)$

출력

첫 번째 줄에 만들 수 있는 수열의 불균형도의 최솟값을 출력한다.

예제 입력 1

3 1
1 2 2

예제 출력 1

예제 입력 2

3 1
2 2 2

예제 출력 2

예제 입력 3

6 3
2 5 7 8 9 11

예제 출력 3

공식 해설 풀이 스포일러 보기 접기

마지막 연산만 원래 수열에 해도 된다

선택한 원소 집합을 $P$, 선택하지 않은 집합을 $C$라 하자. 두 집합 모두 비어 있지 않다. 값이 모두 양수이므로 $C$의 합을 더한 뒤에는 선택한 모든 원소가 선택하지 않은 모든 원소보다 커진다.

따라서 연산 직후의 불균형도는
\[ F_P(A)=\max_{i\in P}A_i+\sum_{i\in C}A_i-\min_{i\in C}A_i \]
이다. 뒤의 두 항은 $C$에서 가장 작은 원소 하나를 제외한 합이다.

고정된 $P$에 대해 이 식은 어느 원소가 증가하더라도 감소하지 않는다. 앞선 연산들은 원소를 증가시키기만 하므로, 여러 번 연산한 결과보다 원래 수열에 마지막 연산만 한 결과가 나쁘지 않다. 따라서 정확히 한 번 연산하는 경우만 보면 된다.

정확히 $K$개를 선택해도 된다

선택한 수가 $K$보다 작으면 $C$에 적어도 두 원소가 남아 있다. 그중 최댓값 $c$ 하나를 $P$로 옮기고 최솟값 하나는 남겨 두자. 이전의 선택 집합 최댓값을 $p$라 하면 식의 변화량은
\[ \max(p,c)-p-c\le0 \]
이다. 따라서 $K$개를 선택할 때까지 늘려도 불균형도는 커지지 않는다.

최솟값과 최댓값의 역할로 선택하기

수열을 $b_1\le\cdots\le b_N$으로 정렬하고 전체 합을 $T$라 하자. 목적함수는
\[ F=T-H,\qquad H=\sum_{i\in P}b_i-\max_{i\in P}b_i+\min_{i\in C}b_i \]
이므로 $H$를 최대화하면 된다.

$K=1$이면 $H$에는 선택하지 않은 최솟값만 남는다. 가장 작은 $b_1$을 선택할 때 이 값이 가장 커져 답은 $T-b_2$이다.

$K=N-1$이면 선택하지 않은 원소가 하나라 그 합과 최솟값이 상쇄된다. 선택한 원소의 최댓값만 최소화하면 되므로 $b_N$을 남겨 두고 나머지를 선택한다. 답은 $b_{N-1}$이다. $N=2$에서는 앞의 경우와 답이 같다.

이제 $2\le K\le N-2$라 하자. $H$는 선택 집합에서 최댓값 하나를 뺀 $K-1$개와 선택하지 않은 최솟값 하나, 총 $K$개 원소의 합이다. 전체 최솟값 $b_1$은 선택했다면 제거되는 최댓값 외의 원소로, 선택하지 않았다면 그 집합의 최솟값으로 포함된다. 전체 최댓값 $b_N$은 선택했다면 제거되는 최댓값이고, 선택하지 않았다면 남은 원소가 적어도 두 개이므로 최솟값으로 고를 필요가 없다. 같은 값이 여러 개면 해당 역할의 원소를 위치로 구분하여 고르면 된다.

따라서 $H$는 $b_1$과 $b_2,\cdots,b_{N-1}$ 중 가장 큰 $K-1$개의 합을 넘지 못한다. 가장 큰 $K$개를 선택하면 바로 이 상한을 달성한다. 이 경우의 답은
\[ b_N+\sum_{i=2}^{N-K}b_i \]
이다.

정렬 후 위 세 경우에 따라 합을 계산한다. 시간복잡도는 $O(N\log N)$, 공간복잡도는 $O(N)$이다. 계산하는 합은 최대 $N\cdot10^9$까지 커질 수 있다.

구현 · 언어별 풀이

C++ 구현

수열을 vector<long long>에 저장하여 정렬하고 전체 합을 구한다. $K=1$을 먼저, 그다음 $K=N-1$을 처리하면 $N=2$도 포함된다. 나머지 경우에는 $0$ 기반 인덱스 $1$부터 $N-K-1$까지의 합에 마지막 원소를 더한다.

C 구현

값을 qsort하고 합은 long long으로 계산한다. K=1 분기를 먼저 검사해야 N=2에서 K=N-1과 겹치는 경우가 올바르다. 일반 경우 가운데 합의 인덱스 범위를 정확히 맞춘다.

Python 구현

정렬 뒤 공통 풀이의 K=1, K=N-1, 나머지 분기를 순서대로 적용한다. 중복 값을 제거하지 않는다. 0-based 가운데 구간을 슬라이스로 표현할 때 끝 인덱스가 제외된다는 점을 확인한다.

Java 구현

int[]를 Arrays.sort하고 합은 long으로 누적한다. K=1을 먼저 처리하고 K=N-1 분기를 뒤에 둔다. int끼리 합한 결과를 마지막에 long으로 바꾸지 않는다.

Rust 구현

Vec<i64>를 sort_unstable()하고 경계 분기를 먼저 처리한다. 가운데 구간의 범위는 끝 제외 슬라이스로 맞춘다. 정렬 후 같은 값도 서로 다른 사람의 값이므로 유지한다.

JavaScript 구현

값 배열을 숫자 비교 함수로 정렬하고 공통 풀이의 세 분기를 적용한다. 합은 최대 $2\cdot10^{14}$라 Number로 정확하며 Int32Array에는 합을 넣지 않는다.

a.sort((x, y) => x - y);
let answer;
if (k === 1) {
    answer = a.reduce((s, x) => s + x, 0) - a[1];
} else if (k === n - 1) {
    answer = a[n - 2];
} else {
    answer = a[n - 1];
    for (let i = 1; i < n - k; i++) answer += a[i];
}
console.log(answer);

$N=2$에서는 첫 분기로도 올바른 답을 얻는다. 가운데 합의 $0$부터 시작하는 인덱스는 $1$부터 n - k - 1까지이다. 값이 같아도 서로 다른 원소이므로 중복을 제거하지 않는다.

제출

언어