問題概要

2 つの数列 $A_{1}, A_{2}, \dots, A_{N}$ と $B_{1}, B_{2}, \dots, B_{M}$ が与えられる。

これらの数列から 1 つずつの値を選ぶ。選んだ値の差が $D$ 以下となるようにすることが可能かどうかを判定し、可能ならば選んだ値の和の最大値を求めよ。

制約

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

考えたこと

この問題は愚直に全探索で解こうと思うと、 $A, B$ から 1 個ずつ選ぶ方法をすべて試す $O(NM)$ の解法となる。これはもちろん間に合わない。

そんな時に有効な考え方は「ある量を固定して考える」というものだ。今回は $A$ から選ぶ要素を固定して考えよう。つまり、

for (int a : A) {
    
}

の中身を埋めようというわけだ。

A の要素 a を固定した問題

$A$ から値 $a$ を選ぶとする。このとき、次のような問題になる。

$B_{1}, B_{2}, \dots, B_{M}$ のうち、 $a$ との値の差が $D$ 以下であるものがあるかを判定し、あるならばそのうちの最大値を求めよ

この処理を実現するためには、C++ ならば upper_bound() を使うのが効果的だ。upper_bound の使い方を次のコードに示す。このコードは、B[i] > v を満たす最小の $i$ を求めるものになっている。計算量は $O(\log M)$ である。

int i = upper_bound(B.begin(), B.end(), v) - B.begin();

なお、upper_bound() を使うためには、 $B$ がソートされている必要があることに注意しよう (あらかじめ最初に $B$ をソートしておけばよい)。

ちなみに、upper_bound() の原理については、次の記事に記した。

qiita.com

upper_bound() を使うと、今回の問題は次のように解ける。

upper_bound() を用いて、B[i] > a + D を満たす最小の $i$ を求める
このとき、は、B[j] <= a + D を満たす最大のとなる
- もし、B[j] >= a - D ならば、a + B[j] が答えである
- そうでなければ、 $a$ との差が $D$ 以下である要素は存在しないと言える

この処理を、各 $a = A_{1}, A_{2}, \dots, A_{N}$ に対して実行すればよい。

全体の計算量は $O((N + M) \log M)$ となる。

コード

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

int main() {
    long long N, M, D;
    cin >> N >> M >> D;
    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());
    
    // a を固定する
    long long res = -1;
    for (long long a : A) {
        int i = upper_bound(B.begin(), B.end(), a + D) - B.begin();
        --i;
        
        if (i < 0) continue;  // これは B[0] > a + D であることを表す
        if (B[i] >= a - D) res = max(res, a + B[i]);
    }
    cout << res << endl;
}

別解：しゃくとり法

$A$ も $B$ もソートしておく。

$A$ を小さい順に見ていきながら、 $B$ の方も小さい順に見ていく。

範囲 $\pm D$ を照らすサーチライトを $A$ 上で動かしながら、それに照らされる物体を $B$ 上で動かしていくようなイメージだ。

具体的には、 $A$ の添字 $i$ と $B$ の添字 $j$ を管理する。最初は $i = j = 0$ としておいて、

$i$ を 1 増やす
B[j] <= a + D を満たすギリギリのところまで $j$ を増やす

という操作を繰り返していく。

計算量は、しゃくとり法自体の部分は $O(N + M)$ だが、 $A, B$ をそれぞれあらかじめソートする部分がボトルネックであり、 $O(N \log N + M \log M)$ となる。

コード

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

int main() {
    long long N, M, D;
    cin >> N >> M >> D;
    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];
    
    // A, B をソートしておく
    sort(A.begin(), A.end());
    sort(B.begin(), B.end());
    
    // しゃくとり法
    long long res = -1;
    int j = 0;
    for (int i = 0; i < N; ++i) {
        // B[j] <= A[i] + D となる最大の j を求める
        while (j < M && B[j] <= A[i] + D) ++j;
        if (j > 0) --j;
        
        if (B[j] >= A[i] - D && B[j] <= A[i] + D) {
            res = max(res, A[i] + B[j]);
        }
    }
    cout << res << endl;
}