Source Code
#include <bits/stdc++.h>
using namespace std;
#define ll long long
int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
#ifndef ONLINE_JUDGE
    freopen("in.txt", "r", stdin);
    freopen("out.txt", "w", stdout);
#endif
    int n, k;
    cin >> n >> k;
    multiset<int> s;
    for (int i = 0, x; i < n; ++i) {
        cin >> x;
        s.insert(x);
    }
    int l = 0;
    int r = n;
    int best = 0;
    while (l <= r) {
        int mid = (l + r) / 2;
        bool flag = 1;
        auto ss = s;
        for (int i = 0; i < k && flag; ++i) {
            for (int j = 0; j < mid && flag; ++j) {
                auto it = ss.find(j);
                if (it == ss.end() || s.empty()) flag = 0;
                else ss.erase(it);
            }
            if (ss.empty() && i + 1 != k) flag = 0;
        }
        if (flag) {
            best = mid;
            l = mid + 1;
        } else r = mid - 1;
    }
    int ans = 0;
    for (int i = 0; i < k; ++i) {
        for (int j = 0; j < best; ++j) {
            s.erase(s.find(j));
        }
        ans += best;
    }
    for (int i = 0; i < k && s.size(); ++i) {
        int mex = best;
        for (; s.size() > k - i - 1;) {
            auto it = s.find(mex);
            if (it == s.end()) break;
            s.erase(it);
            mex++;
        }
        ans -= best;
        ans += mex;
    }
    cout << ans << '\n';
    return 0;
}
Copy
Bitar The Handy Man Heartbeat
GNU G++17
379 ms
10.4 MB
Accepted