ベスパリブ

プログラミングを主とした日記・備忘録です。ベスパ持ってないです。

ABC 181: E問題の解説 (C++)

問題:AtCoder Beginner Contest 181: E - Transformable Teacher

解説

 H Wは事前にソートしておきます。

先生と組む生徒以外の人のペアの身長差は、ソートしたときに隣り合う人でペアを作り続ければよさそう。

図にすれば以下。

黒丸が先生と組む生徒 i、白丸はそれ以外の生徒。

f:id:takeg:20201114161152j:plain
ペアの作りかた

図を眺めると、「先生と組む生徒」以外のペアの身長差の合計 totalは、最初の1回目( i=0のとき)に totalを計算するときだけ O(N)かかりますが、2回目以降( i=1以降)は、しゃくとり法を使えば totalの計算は O(1)で計算できそうです。

あとは「先生と組む生徒」と「先生」の身長差を高速で求めたいですが、これも先生が変身できる身長配列 Wに対して、生徒の身長がどこに入るかを二分探索すればO(logN)で求まります。

これを各生徒iのどの生徒のときに最小になるのかN回探索するので、全体としてはO(N logN)で計算できます。

実装

C++での解法です。

#include <bits/stdc++.h>
#define _USE_MATH_DEFINES  // M_PI等のフラグ
#define MOD 1000000007
#define COUNTOF(array) (sizeof(array)/sizeof(array[0]))
#define rep(i,n) for (int i = 0; i < (n); ++i)
using namespace std;
using ll = long long;
using P = pair<int,int>;


void solve(){
    int N, M;
    cin >> N >> M;
    vector<int> H(N);
    vector<int> W(M);
    for (int i=0; i<N; i++) cin >> H[i];
    for (int i=0; i<M; i++) cin >> W[i];
    sort(H.begin(), H.end());
    sort(W.begin(), W.end());

    /*
    i=0から順に、先生と組む場合の身長差を計算をする
    i=1のときは、しゃくとり法を使えば先生と組む以外の人たちの身長差は一瞬で計算できるO(1)
    先生がどの身長になればよいかは、二分探索で求めればO(logN)
    よって、この計算量はO(NlogN)でいけそう
    */
    vector<int> diffT(N,0);  // 生徒iと、先生との最小の身長差
    for (int i=0; i<N; i++) {
        auto it = lower_bound(W.begin(), W.end(), H[i]);
        if (it == W.begin()) {
            diffT[i] = W[0]-H[i];
        }
        else if (it == W.end()) {
            diffT[i] = H[i]-W.back();
        }
        else {
            diffT[i] = min(abs(H[i]-*it), abs(H[i]-*(it-1)));
        }
    }

    // i=0の生徒が先生とペアになるとき
    int total = 0;
    for (int i=1; 2*i<N; i++) {
        total += H[2*i]-H[2*i-1];
    }
    int ans = total + diffT[0];

    // i=1~N-1の生徒が先生とペアになるときを、しゃくとり法で求める
    for (int i=1; i<N; i++) {
        if (i%2==1) {
            total -= H[i+1] - H[i];
            total += H[i+1] - H[i-1];
            ans = min(ans, total+diffT[i]);
        }
        else {
            total -= H[i] - H[i-2];
            total += H[i-1] - H[i-2];
            ans = min(ans, total+diffT[i]);
        }
    }

    // 出力
    cout << ans << endl;
}


int main(int argc, char const *argv[]){
    solve();
    return 0;
}

計算量は、ソートした処理も含めると O(N logN + M log M)です 。