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

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

AtCoder ABC 119 D - Lazy Faith (水色, 400 点)

二分探索で lower_bound とかきっちり無意識的に使いこなせるようになりたいのんな!!!

問題へのリンク

問題概要

一次元の世界を考える。
A 個の神社と、B 個の寺が並んでいる (各神社と各寺の位置情報が 1 つの整数値で与えられる)。

以下の Q 個のクエリに答えよ:

  • 座標値 x が与えられる。x から出発して「神社を 1 個以上」「寺を 1 個以上」訪れるのに必要な総移動距離の最小値を求めよ

制約

  •  1 \le A, B, Q \le 10^{5}

考えたこと

クエリが  Q 個あるので、それぞれのクエリに爆速で答えなければならない。 O(\log) のオーダーで答えたいところである。

さて、最適解は以下のいずれかの形であることはわかる

  • まずいずれかの「神社」を訪れて、そこから最も近い「寺」に訪れる
  • まずいずれかの「寺」を訪れて、そこから最も近い「神社」に訪れる

ここで、1 個目の場合について、最初に訪れる神社が必ずしも初期位置から最も近い位置がよいわけではないことに注意する。なぜなら、初期位置から最も近い神社に行っても、そこから最寄りの寺がすごく遠いかもしれない!!!!!!!!

しかしそれでも、

  • 初期位置から前方の最寄りの神社
  • 初期位置から後方の最寄りの神社

以外を狙う意味はないことがわかる。何故なら、これら以外の神社に行こうと思うと、「前方最寄り」「後方最寄り」のうちのいずれか一方には既に通りかかることになるからだ。

というわけで初回の行動は以下の 4 パターンで、2 回目の行動はそこから最寄りの「寺」か「神社」に行けばよい。

  • 最初に前方最寄りの神社に行く
  • 最初に後方最寄りの神社に行く
  • 最初に前方最寄りの寺に行く
  • 最初に後方最寄りの寺に行く

前方最寄りや後方最寄りを求める方法

残る問題は、

  • 現在位置  x から出発して
  • 前方最寄り、または後方最寄りの、神社なり寺なりに行く方法

を素早く求めること。これは二分探索を用いることで実現できる。神社なり寺なりの位置情報が vector v (ソート済み) で与えられているとして、以下のようにもとめられる。

前方の index (自分自身を含んで良い場合)

upper_bound(v.begin(), v.end(), x) - v.begin() - 1;

前方の index (自分自身を含まない場合)

lower_bound(v.begin(), v.end(), x) - v.begin() - 1;

後方の index (自分自身を含んで良い場合)

lower_bound(v.begin(), v.end(), x) - v.begin();

後方の index (自分自身を含まない場合)

upper_bound(v.begin(), v.end(), x) - v.begin();

注意点として、それぞれ配列外アクセスになる可能性があるので、予め番兵として

  • -INF, INF の位置に神社や寺があるとしておく

とするのがよかったりする。

#include <iostream>
#include <vector>
#include <cmath>
#include <algorithm>
using namespace std;
void chmin(long long &a, long long b) { if (a > b) a = b; }
const long long INF = 1LL<<58;

// vector v において、x の前方と後方の index を求める
template<class T> int former(const vector<T> &v, T x) {
    return upper_bound(v.begin(), v.end(), x) - v.begin() - 1;
}
template<class T> int latter(const vector<T> &v, T x) {
    return lower_bound(v.begin(), v.end(), x) - v.begin();
}

int main() {
    int A, B, Q; cin >> A >> B >> Q;
    vector<long long> s(A), t(B);
    for (int i = 0; i < A; ++i) cin >> s[i];
    for (int i = 0; i < B; ++i) cin >> t[i];
    s.push_back(INF); s.push_back(-INF); sort(s.begin(), s.end());
    t.push_back(INF); t.push_back(-INF); sort(t.begin(), t.end());
    
    for (int _ = 0; _ < Q; ++_) {
        long long x; cin >> x;
        long long res = INF;
        
        // 最初に s (前方に行くか後方に行くか場合分け)
        for (int i = 0; i < 2; ++i) {
            long long first = (i ? s[former(s, x)] : s[latter(s, x)]);
            // 次に t (現在地の前後を見る)
            for (int j = 0; j < 2; ++j) {
                long long second = (j ? t[former(t, first)] : t[latter(t, first)]);
                chmin(res, abs(x - first) + abs(first - second));
            }
        }
        // 最初に t (前方に行くか後方に行くか場合分け)
        for (int i = 0; i < 2; ++i) {
            long long first = (i ? t[former(t, x)] : t[latter(t, x)]);
            // 次に t (現在地の前後を見る)
            for (int j = 0; j < 2; ++j) {
                long long second = (j ? s[former(s, first)] : s[latter(s, first)]);
                chmin(res, abs(x - first) + abs(first - second));
            }
        }
        cout << res << endl;
    }
}