AtCoder ABC 321 D - Set Menu (緑色, 400 点)

「二分探索 lower_bound()」「累積和」を活用する、とてもとても典型的かつ教育的な問題ですね。

• 問題へのリンク

問題概要

サイズ の数列 と、サイズ の数列 が与えられる。

これらの数列から 1 個ずつ選んでできる 通りの各ペア についての、 の総和を求めよ。

制約

• 

考えたこと

単純に全探索すると の計算量となって間に合わない。この手の問題で考えることは、いつも決まっている。それは

---

数列 の側から選ぶ要素 を固定したときの、数列 全体について調べた値を高速に求められないだろうか

---

このように、ある値を固定して考えるというのは、茶色から緑色へと上がっていく上で必須級のスキルとなる。

数列 の要素を固定して考える

具体例で考えてみよう。

• 
• 
• 

とする。ここで、数列 からは要素 を固定することにする。このとき、 の各要素は

• 未満の値はそのまま足す
• 以上の値はすべて 9 となる

というようになっていることがわかる。下図のようになる (前者は青色、後者は赤色)。

この図のうち、

• 青色の部分については、累積和を用いると高速に求められる
• 赤色の部分については、 × (その個数) によって高速に求められる (A の分も足すので となる)

ということが分かる。数列 を小さい順に並べたとき、どこから 以上の値になるかについては、関数 lower_bound() を用いた二分探索法のよって求められる。lower_bound() については、詳しくは次の記事にて。

drken1215.hatenablog.com

以上の解法の計算量は

• 最初に数列 を小さい順にソートしておく ：
• 数列 の累積和を求めておく：
• 数列 の各要素についての数列 の処理には を要するので、数列 全体の要素についての処理は：

ということで、全体で の計算量となる。

コード

#include <bits/stdc++.h>
using namespace std;

int main() {
    // 入力
    long long N, M, P;
    cin >> N >> M >> P;
    vector<long long> A(N), B(M);
    for (int i = 0; i < N; ++i) cin >> A[i];
    for (int i = 0; i < M; ++i) cin >> B[i];
    
    // B を小さい順に並び替えて、累積和も求める
    sort(B.begin(), B.end());
    vector<long long> S(M+1, 0);
    for (int i = 0; i < M; ++i) S[i+1] = S[i] + B[i];
    
    // 数列 A の各要素について処理する
    long long res = 0;
    for (long long a : A) {
        // B[it] >= P - a となる最小の it を求める
        int it = lower_bound(B.begin(), B.end(), P - a) - B.begin();
        
        // 前半と後半に分けて求める (A の分も足すのを忘れずに)
        res += (S[it] + a * it) + P * (M - it);
    }
    cout << res << endl;
}