Official

E - How to Win the Election Editorial by MtSaka


この問題は二分探索の練習問題となっています。

まず、\(N=M\) の場合は全候補者が当選します。以降\(M<N\) であるとします。また、候補者 \(i\)\(i\) 番目の候補者を指します。

以下の判定問題を解くことができれば二分探索を用いてこの問題の答えを求めることができます。

すべての開票が終わった時点で候補者 \(i\) が合計で \(X\) 票獲得したとした時に自分以外の上位 \(M\) 人の候補者がそれぞれ合計 \(X+1\) 票以上獲得していることは可能か。

以上の判定問題において可能である場合は、 \(i\) 番目の候補者は合計 \(X\) 票の時に、当選しない可能性があるため当選を確定させることができません。逆に、可能でない場合は \(i\) 番目の候補者は合計 \(X\) 票獲得することができれば確実に当選します。

実際にこの判定問題を高速に解くことを考えます。 まず、簡単のため \(A\) が昇順であるとします。

すべての開票が終わった時点で、候補者 \(i\) 以外の上位 \(M\) 人がそれぞれ \(X+1\) 票以上獲得するのに合計で追加票が最小で何票必要かを求めたいです。 この値を \(C\) としたとき、\(C>K-\sum_{i=0}^N A_i\) を満たす場合は判定問題は不可能となり、そうでない場合は可能になります。

この上位 \(M\) 人を現在すでに開票されている時点での候補者 \(I\) 以外の上位 \(M\) 人と固定した時にこの最小値は達成されます。 \(1 \leq i \leq N-M\) の場合は現時点での上位 \(M\) 人は 候補者 \(N-M+1,\ldots ,N\) です、\(N-M+1 \leq i \leq N\)の場合は候補者 \(N-M,N-M+1,\ldots,i-1,i+1,\ldots,N\) が上位 \(M\) 人になります。 これらの候補者全員が \(X+1\) 票以上獲得するにはあと何票必要かを求めます。

これらの候補者のうちすでに \(X+1\) 票以上獲得している場合は追加の票は必要ありませんが、\(j( \leq X)\) 票であった場合は \(X+1-j\) 票追加で必要です。この値を各候補者について愚直に加算していくという方法で求められます。この時判定問題を解く時間計算量は \(O(M)\) となります。しかし、このままでは十分に高速ではありません。

高速化するには、上位 \(M\) 人が基本的に連続する番号の人たちであること、\(X+1\) 票以上すでに獲得しているかは \(j\) について単調であることから累積和と二分探索を組み合わせることで判定問題を解くのに求めるべき値を求めることができます。 具体的には、二分探索により \(X+1\) 票未満しか獲得できてない候補者の番号の最大値を求めて、その番号以下の上位 \(M\) 人の候補者の人数を求めて 人数 \(\times (X+1)\)を加算します。そこから、先ほど考慮した人全員の現在獲得している票の総和を引きます。 最後に、候補者 \(i\)\(X\) 票獲得するのに必要な追加票を加算し、出た値が必要な追加票の最小値です。 考慮する上位 \(M\) 人の番号が連続でない場合があるため、実装の際は注意が必要です。詳細は実装例を参照してください。

よって、各候補者に対して、この二分探索を \(O(\log K)\) 回行い、判定問題は \(O(\log N)\) で解くことができるため、全体で時間計算量 \(O(N \log K \log N)\) で解くことができます。

実装例(C++):

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
int main() {
    int n, m;
    ll k;
    cin >> n >> m >> k;
    vector<ll> a(n);
    for (auto& e : a) cin >> e;
    const ll rem = k - accumulate(a.begin(), a.end(), 0LL);

    if (n == m) {
        for (int i = 0; i < n; ++i) cout << 0 << " \n"[i == n - 1];
        return 0;
    }

    vector<int> ord(n);
    iota(ord.begin(), ord.end(), 0);
    sort(ord.begin(), ord.end(), [&](int i, int j) { return a[i] < a[j]; });
    auto b = a;
    sort(b.begin(), b.end());
    vector<ll> sumb(n + 1);  // 昇順にしたAの累積和
    for (int i = 0; i < n; ++i) sumb[i + 1] = sumb[i] + b[i];

    vector<ll> ans(n, -1);
    for (int i = 0; i < n; ++i) {
        ll l = -1, r = rem + 1;
        while (r - l > 1) {
            ll mid = (l + r) / 2;
            ll rid = lower_bound(b.begin(), b.end(), b[i] + mid + 1) - b.begin();  // b[rid - 1]<b[i]+mid+1 となる最大のridを求めている
            ll lid = n - m - (i >= n - m ? 1 : 0);  // i<n-mの時は 候補者n-m,...,nが上位 m 人でi >=n-m の時は候補者n-m-1,...,i-1,i+1,..,.n が上位 m 人になる
            ll cnt = 0;
            if (rid > lid) cnt += (rid - lid) * (b[i] + mid + 1) - (sumb[rid] - sumb[lid]);
            if (lid <= i && i < rid)
                cnt--;  // この時、すでに候補者iが b[i]+mid+1 票獲得していることにしているので 1を引く
            else
                cnt += mid;  // 候補者 i の必要な票を加算する
            if (cnt > rem)
                r = mid;
            else
                l = mid;
        }
        if (l == rem)
            ans[ord[i]] = -1;
        else
            ans[ord[i]] = r;
    }

    for (int i = 0; i < n; ++i) cout << ans[i] << " \n"[i == n - 1];
}

posted:
last update: