ABC_216_E - Amusement Park

[zeikar]

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;
}