文学徒が競プロをするブログ

競技プログラミングの精進を記録していきます

ABC149E: Handshake

いい問題でした。The 500点って感じ。

【問題概要】

N個の要素を持つ数列 Aが与えられる。
\sum_{x=1}^{n} \sum_{y=1}^{n} A_x+A_y とし、全ての組み合わせにおいて A_x+A_yを降順に並べた時、 M番目までの総和を求めよという問題。
かなり汎用性が高そうな題材なので、無限回既出なんだろうなという気持ちになる。

N^2が最大で 10^{10} になるので、全列挙では残念ながら間に合わない。どうにか N^2かけずに組み合わせを考えたい気持ちになる。
解法を先に提示しておく。

  1. A_x+A_yを降順に並べた際、 M番目に来る数を求める
  2. 操作1で求めた M番目の数を Kと置いた時、 K \leqq A_x+A_yとなる数の総和を求める

（最後に調整が必要な場合もあるが）実はこの2つの操作さえ処理できればこの問題を解くことが出来る。そしてこの処理は一見 N^2かかってしまいそうに見えるが、前者は O(N\log{}^2N)、後者は O(N\log{}N)で処理することが可能である。
以下、各操作について述べる。

1. A_x+A_yを降順に並べた際、 M番目に来る数を求める

本質を言う。
真偽判定として「 M番目に来る数を仮にKとする。 K \leqq A_x+A_yとなるような組み合わせの数が本当に M個以上あるか？」を考え、二分探索をする。
ここがわからない人に、この問題は早いかも知れない。これなんかをやってまず二分探索を理解するといいと思う。

さて、ではどういったアルゴリズムを用いるかだが、 K \leqq A_x+A_yとなる組み合わせの数が求められれば上記の真偽判定が出来ることになる。ここでは get\_count(K) := K \leqq A_x+A_yとなる組み合わせの数、という関数を考えてみる（はまやんさんリスペクト）。
さて、ここでこの不等式を移項すると K-A_x \leqq A_yとなる。これは各 A_xについて、 K-A_xよりも大きい A_yの数の合計が、 K \leqq A_x+A_yの組み合わせの数と同値であることを意味している。

要は A_xを1つ固定すると、 K \leqq A_x+A_yにするために必要な A_yの個数が A_xそれぞれについてわかるということである。
で、これはAを昇順ソートしておくとlower_boundで求められる。詳細な雰囲気は以下のコードで察してほしい。

ll get_count(ll K, const vector<ll> &a) {
    // aはソート済み
    int n = (int)a.size();
    ll ret = 0ll;
    rep(i,n) {
        // K-Axをbと置く
        ll b = K-a[i];
        auto itr = lower_bound(a.begin(), a.end(), b);
        // lower_boundで得られたイテレータを終端のイテレータから引くと
        // K-Ax以上であるAyの個数が分かるので、それを各Axについて足していく
        ret += (a.end()-itr);
    }
    return ret;
}

Kを二分探索する準備ができたので、やる。

// binary search
ll ok = 0; ll ng = LINF;
while (abs(ok-ng)>1) {
    ll mid = (ok+ng)/2;
    ll cnt = get_count(mid,a);
    if (cnt>=m) ok = mid;
    else ng = mid;
}

ok という変数に M番目に来る数 Kが入っている。
これで A_x+A_yを降順に並べた際、 M番目に来る数を求めることができた。

2. K \leqq A_x+A_yとなる数の総和を求める

get\_sum(K) := K \leqq A_x+A_yとなる数の総和、という関数を考えてみる（はまやんさんリスペクトその2）。
やっぱりここでも K-A_x \leqq A_yと置く。 A_xを1つ固定すると、 K \leqq A_x+A_yにするためには K-A_x以上の A_y全てを A_xと足し合わせていけばいい。
で、やっぱりAを昇順ソートしておくとlower_boundと累積和（ A_yを全て足していると間に合わないので）で求められる。詳細な雰囲気は以下のコードで察してほしい。

// s := Aの総和
// sum := Aの累積和
ll get_sum(ll K, ll s, const vector<ll> &a, const vector<ll> &sum) {
    ll ret = 0;
    int n = a.size();
    rep(i,n) {
        // K-Axをbと置く
        ll b = K-a[i];
        auto itr = lower_bound(a.begin(), a.end(), b);
        // lower_boundで得られたイテレータで始点のイテレータを引くと
        // AyがK-Ax以下となるインデックスがわかるので
        ll idx = itr-a.begin();
        // Aの総和からそこまでの累積和を引くと
        // K-Ax以上のAyの総和が分かる
        ret += (s-sum[idx]);
        // cnt := 足したAyの個数 
        ll cnt = n-idx;
        // Axをその分増やす
        ret += (a[i] * cnt);
    }
    return ret;
}

これでできた。

3. 足しすぎた分の調整

最初のサンプルはここまでやってきたコードで通すことが出来る。
しかし最初のサンプル

5 3
10 14 19 34 33

を

5 2
10 14 19 34 33

にしても同じ答えが出てしまう。
ここで \sum_{x=1}^{n} \sum_{y=1}^{n} A_x+A_y とし、全ての組み合わせにおいて A_x+A_yを降順に並べたものを見てみると
{68, 67, 67, 66, 53, 53, 52, 52, .... } となっているのがわかる。
上から3つ目も2つ目も67になってしまっているため、 M=3の答えが M=2の時にも出てしまうのである。
だから、この分を引いてやればいい。
get_count(K)で同じ数になっている中では1番右の数（1番インデックスが大きい数）が得られるので、そこから Mを引いた個数分Kを引いてやる。

ll K = ok;
ll ans = get_sum(K,s,a,sum);
ans -= (get_count(K,a)-m)*K;

これで全部通る。
ACコードはこれ。