ABC_216_E - Amusement Park #31
Description
Problem link
https://atcoder.jp/contests/abc216/tasks/abc216_e
Problem Summary
n개의 a 배열이 주어지고 최대 k개를 골라서 최대가 되도록 고르는 문제.
단, 한번 고를 때마다 a 배열의 원소의 값은 1씩 감소한다.
Solution
뭔가 이분 탐색으로 가능할 것 같은데 생각보다 쉽지 않았다.
먼저 k개를 넘지 않는 최대 개수를 고르는 지점을 이분 탐색을 이용해서 구할 수 있다.
k 개에서 남은 개수만큼 더 고를 수 있는데 이건 따로 한 번 더 탐색해서 구하면 된다.
예제1로 보면 이분 탐색으로 k개를 넘지 않게 고를 수 있는 값인 100을 구함 (4개 선택 가능)
k 가 5이므로 1개를 더 추가로 고를 수 있는데 이건 아래쪽 for 문으로 cnt를 구해 k 에서 뺀 만큼 99 를 곱해서 더하도록 했다.
Source Code
#include <iostream>
using namespace std;
long long a[100000];
int main() {
long long n, k;
cin >> n >> k;
long long high = 0;
long long low = 0;
long long ans = 0;
for (int i = 0; i < n; i++)
{
cin >> a[i];
high = max(high, a[i]);
}
while (low <= high)
{
long long mid = (high + low) / 2;
long long cnt = 0;
for (int i = 0; i < n; i++)
{
if (a[i] < mid)
{
continue;
}
long long n = a[i] - mid + 1;
cnt += n;
}
if (cnt > k)
{
low = mid + 1;
}
else
{
high = mid - 1;
}
}
long long cnt = 0;
for (int i = 0; i < n; i++)
{
if (a[i] < low)
{
continue;
}
long long n = a[i] - low + 1;
cnt += n;
ans += n * (a[i] + low) / 2;
}
ans += max((k - cnt) * (low - 1), 0LL);
cout << ans << endl;
}