[AtCoder] ABC 149 E – Handshake

2019年12月31日2020年12月17日

問題

E - Handshake

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly progr ...

https://atcoder.jp/contests/abc149/tasks/abc149_e

方針

一般ゲストの並びを変えても影響がないので、降順に並び替えます。つまり、\( A_{i}\geq A_{i + 1} \ (0 \leq i \leq N – 2)\) を満たすとします。

最終的な幸福度が最大値のとき、\( 1 \) 回の握手で増加する幸福度の最小値 \( x \)

最終的な幸福度が最大値を取るとき、\( 1 \) 回の握手で増加する最大値は \( 2A_0 \) ですが、ここでは\( 1 \) 回の握手で増加する幸福度の最小値 \( x \) を求めます。ここで、\( d(i) \) をパワーが \( i \) 以上の人数とします。これは、累積和を利用することで求めることができます。また、\( d(1) = N \) であることは明らかです。次に、\( x \) を二分探索によって求めます。\( x \ ( 1 \leq x \leq 2 \times 10^5)\) を固定したとき、\( x \) 以上を満たす握手のパターン \( s \) を求めます。これは、

\[ s = \sum_{k = 0}^{N – 1} d(\max(0, \ x – A_k)) \]

として計算できます。\( s \geq M \) を満たすとき、\( M \) 回の握手の中で、\( 1 \) 回の握手で \( x \) 以上の幸福度が得られることが保証されます。これは、最終的な幸福度の最大値が \( Mx \) 以上であり、少なくとも \( 1 \) 回以上 \( x \) を満たす握手が存在することが分かります。

握手のシミュレーション

\( x \) 以上の握手のパターンの中で、上位 \( M \) 個の握手を行います。ここで、\( x + 1 \) 以上を満たす握手のパターンを先に行います。\( A_i \) の累積和 \( b_i \) を次のように定義します。\( b_0 = 0\) として、

\[ b_{i} = \sum_{k = 0}^{i – 1} A_k \ (1 \leq i \leq N )\]

とします。\( x + 1 \) 以上を満たす握手のパターンを行ったときの増加する幸福を \( f \) とします。これは、次のようにして求めることができます。ゲスト \( i \) と握手する人数を \( t_i \) とし、既に握手した人数を \( k \) とします。初期値は、\( k = 0 \) です。

\( v_i = \max(0, \ x + 1 – A_i) \)
\( t_i = \max(0, \ \min(d(v_i), M – k)) \)
\( f \leftarrow f + b(t_i) + t_i \times A_i \)
\( k \leftarrow k + t_i \)
\( i \leftarrow i + 1\) 1. に戻る

したがって、答えは、\( f + x(M – k) \) となります。

また、最終的な \( k \) の値は \( k < M \) となります。これは、\( x + 1 \) 以上の握手のパターンが \( M \) 未満であり、\( x \) 以上のパターンが \( M \) 以上であるためです。

コード

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int main() {
    ll N, M;
    cin >> N >> M;
    vector<ll> A(N);
    for (int i = 0; i < N; i++) {
        cin >> A[i];
    }
    sort(A.begin(), A.end(), greater<ll>());
    ll b[N + 1]{};
    for (int i = 0; i < N; i++) {
        b[i + 1] = b[i] + A[i];
    }
    ll A_MAX = 2 * 100000 + 1;
    ll d[A_MAX]{}; // d[i]: パワーi以上を持つ人の数
    for (int i = 0; i < N; i++) {
        d[A[i]]++;
    }
    for (int i = A_MAX - 1; i >= 1; i--) {
        d[i - 1] += d[i];
    }
    ll l = -1;
    ll r = A_MAX;
    while (r - l > 1) {
        ll m = (l + r) / 2;
        ll s = 0;
        for (int i = 0; i < N; i++) {
            s += d[max(1ll, m - A[i])];
        }
        if (s >= M) {
            l = m;
        } else {
            r = m;
        }
    }
    
    ll ans = 0;
    ll s = 0; // 握手した人の合計 
    for (int i = 0; i < N; i++) {
        int v = max(0, (int)(r - A[i]));
        ll t = max(0ll, min(d[v], M - s));
        ans += b[t] + t * A[i];
        s += t;
    }
    ans += l * (M - s);
    
    cout << ans << "\n";
    
    return 0;
}