けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

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

AtCoder AtCoder400点 ABC-D 緑色diff そのまま覚えたい典型問題 典型要素を詰め合わせた教育的問題 ある量を固定して考える 数列 2つの数列 累積和 【問題集】累積和・二分探索法・しゃくとり法 二分探索:lower_bound 二分探索 解空間:O(N^2)個のペア 解空間:O(N^2)通りの選択肢 総和を求める

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

問題概要

サイズ N の数列 A_{1}, A_{2}, \dots, A_{N} と、サイズ M の数列 B_{1}, B_{2}, \dots, B_{M} が与えられる。

これらの数列から 1 個ずつ選んでできる NM 通りの各ペア (a, b) についての、 \min(a+b, P) の総和を求めよ。

制約

  • 1 \le N, M \le 2 \times 10^{5}

考えたこと

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

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

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

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

具体例で考えてみよう。

  • A = (3, 5, 7)
  • B = (1, 3, 6, 8, 10, 11, 12)
  • P = 14

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

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

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

この図のうち、

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

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

drken1215.hatenablog.com

以上の解法の計算量は

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

ということで、全体で O((M + N) \log M) の計算量となる。

コード

#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;
}