算法
(二分) $O(N\log 2e9)$
本题其实就是问
从序列 $B = (1, 2, \cdots, A_1, 1, 2, \cdots, A_2, \cdots, 1, 2, \cdots, A_N)$ 中至多选 $K$ 个元素,总和的最大值是多少?
二分,考虑以下决策:
序列 $B$ 中是否存在大于 $X$ 的元素个数小于 $K$?$(*)$
$(*)$ 对于 $X = -1$ 时为 false
(左边界设置为 -1
是因为存在 $K$ 超过序列 $B$ 的长度的情况),$X=2e9$ 时为 true
记 f(X)
表示大于 $X$ 的元素个数,$f(X) < K \leqslant f(X - 1)$
因此,我们可以二分 $X$ 找到满足 $(*)$ 的最小值 $X$,不妨记为 $M$
当从序列 $B$ 中选择元素时
- 大于 $M$ 的元素全都选
- 剩下 $K - f(M)$ 个都选择 $M$
C++ 代码
#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); ++i)
using std::cin;
using std::cout;
using std::max;
using std::vector;
using ll = long long;
int main() {
int n, k;
cin >> n >> k;
vector<int> a(n);
rep(i, n) cin >> a[i];
ll wa = -1, ac = 2e9;
auto f = [&](ll x) {
ll res = 0;
rep(i, n) res += max(0ll, a[i] - x);
return res;
};
while (wa + 1 < ac) {
int wj = (wa + ac) / 2;
if (f(wj) < k) ac = wj; else wa = wj;
}
ll ans = 0;
auto asum = [&](ll l, ll r) {
return (l + r) * (r - l + 1) / 2;
};
rep(i, n) {
if (a[i] <= ac) continue;
ans += asum(ac + 1, a[i]);
}
ans += ac * (k - f(ac));
cout << ans << '\n';
return 0;
}