ABC149 E - Handshake

ABC149 E - Handshake

問題へのリンク

問題概要

i,jをM回選んだとき、A[i] + A[j]の最大値を答えよ。ただし、同じi,jを選べない。

解説

方針として、全てのi,jを列挙して大きい方からM個とってくれば良いことがわかる。 ただし、この方針では計算量がO(N²)となる。

Aをソートし、i,jの対応表を考えたとき、大きい値は右下によっていることを想像できる。もっと厳密に言えば、あるiについて二分探索をすることを考えれば、O(NlogN)で和がX以下（以上）のものを列挙できる。これを「二次元における二分探索」という典型として持っておくと嬉しそう。今回、ある閾値を設けて、それ以上の値がいくつあるかを数えることができる。また、この閾値と個数には単調性がある。閾値を大きくすると、それ以上の値は少なくなる。つまり、どの閾値ならギリギリM個以下になるのかということを、二分探索を用いて高速に計算可能なのである。

ある閾値X以上だとM個以下で済むが、X未満だとM個を超えてしまうという境界を探索する。Xに設定したとき、Mちょうどならいいが、そうとも限らない。X–1にするとMを超えるが、XだとMより少ないというなんとも中途半端な結果が返ってくることがある。このとき、どうすれば良いかというと、X以上の値を列挙したあと、Mより少ない分だけX-1を足せるということになるので、計算の途中でいくつ使ったかを計算しておき、足りない分を後で合わせるということをすればよい。

提出

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(20);

    ll N, M;
    cin >> N >> M;
    vector<ll> A(N);
    rep(i,N)cin>>A[i];
    sort(all(A));
    // X以上のものだけを使ったときの個数がM以下か
    auto check = [&](ll X){
        ll cnt = 0;
        rep(i,N){
            int add = A.end() - lower_bound(all(A), X - A[i]);
            cnt += add;
        }
        return cnt < M;
    };
    ll ng = 0;
    ll ok = INFLL;
    while(ok - ng > 1){
        ll mid = (ok + ng) / 2;
        if(check(mid)){
            ok = mid;
        }
        else{
            ng = mid;
        }
    }
    cerr << ng << " " << ok << endl;
    //累積和で高速化パート
    vec<ll> cum(N+1, 0);
    rep(i, N)cum[i+1] = cum[i] + A[i];
    //ok以上のものだけ使えばM個以内に収まる
    //ngを何個か入れることを考える
    //ngが入る余地があるときがある
    //ok以上でM以下だが、Mピッタリとは限らない
    //ngにすると超えるということである。
    //余りがあるはずなので、それを加えよう。
    ll ans = 0;
    rep(i,N){
        int idx = lower_bound(all(A), ok - A[i]) - A.begin();
        int cnt = N - idx;
        ans += cum[N] - cum[idx] + A[i] * cnt;
        M -= cnt;
    }
    ans += M * (ok - 1);
    cout << ans << endl;
}