AtCoder ABC 373 E - How to Win the Election (2D, 水色, 500 点)

方針を立てるのは難しくないけど、とにかく重たい問題だった。

• 問題へのリンク

問題概要

人 に対して合計 票が集まった。これらの人のうち条件「自分よりも多くの票を集めた人が 人未満」を満たす人が当選となる。

票のうち途中まで開票されて、各人 には 票が集まった。

各人 について、条件「あと 票を追加すれば当選確実」を満たすような最小の整数 を求めよ (存在しない場合は -1)。

制約

• 
• 

考えたこと

とする (残り票数)。

コンテスト本番では、 の場合に処理がおかしくなっていることに気づかず、デバッグに時間がかかってしまった。 の場合は、全員が当選確実 (自分より票を集めるのは 人以下のため) なので、0 を 個出力すればよい。以降、 としよう。

さて、 がある値で条件を満たすならば、それより大きい値でも当然条件を満たす。このことから、その境界を求めるのに二分探索が使える。つまり、次の判定問題を解けば良い。この判定問題の答えが "Yes" となる最小の を求めれば良いのだ。

---

のうち、 を除いて大きい順に 個を考える。

それらの値を 以上にするのに必要な得票数を とする。

このとき、 であるかどうかを判定せよ。

---

あとは、 さえ求められればよい。これは lower_bound() や累積和を用いて高速に求められる。

全体として、計算量は となる。

コード

#include <bits/stdc++.h>
using namespace std;
using pll = pair<long long, long long>;

int main() {
    long long N, M, K;
    cin >> N >> M >> K;
    vector<long long> A(N), S(N+1, 0);
    vector<pll> pA(N);
    for (int i = 0; i < N; i++) cin >> A[i], pA[i] = pll(A[i], i);
    sort(A.begin(), A.end()), sort(pA.begin(), pA.end());
    for (int i = 0; i < N; i++) S[i+1] = S[i] + A[i];
    long long rem = K - S.back();  // 残り票数

    // 例外処理
    if (M == N) {
        for (int i = 0; i < N; i++) cout << 0 << " ";
        cout << endl;
        return 0;
    }

    vector<long long> res(N);
    for (int i = 0; i < N; i++) {
        int person = pA[i].second;
        long long low = -1, high = rem+1;
        while (high - low > 1) {
            long long x = (low + high) / 2;

            // 上位 num 人の得票数を A[i] + x + 1 以上にするのに必要な票数 sum を求める
            long long sum = 0, border = A[i] + x + 1;
            long long mid = lower_bound(A.begin(), A.end(), border) - A.begin();

            // もし上位 M 人の中に人 person が含まれる場合は M+1 人に広げる
            long long num = M, sub = 0;
            if (i >= N - M) num++, sub = border - A[i];
            long long left = N - num;
            if (mid >= left) sum = border * (mid - left) - (S[mid] - S[left]) - sub;
            if (sum <= rem - x) low = x;
            else high = x;
        }
        res[person] = (high <= rem ? high : -1);
    }
    for (auto val : res) cout << val << " ";
    cout << endl;
}