ABC 181 | E - Transformable Teacher

2 min read読了の目安(約2100字

問題

https://atcoder.jp/contests/abc181/tasks/abc181_e

考えたこと

身長差の合計を最小化するには、とある児童と比べて身長が次に低い児童または先生か、次に高い児童または先生をその児童の隣に選べば良い。
よって児童を昇順でソートし、どの身長W_iでどこに先生を入れればいいか最適解が見つかればそれが答えになる。

以下は入力例1の例である。最適な先生の候補を赤丸のどこかに入れるのが最適解である。

そのままやると先生の候補がM個、赤丸の候補がN+1個、身長の差の合計を求めるのに(N+1)/2回計算が必要になりO(MN^2)となり間に合わない。

先生は入力例1だと以下の赤の場所のとこかに入れる。
ソートされているのでどこに入れるかは二分探索を使いO(logN)で求められる。

なお、どこに入れるかのindexは以下のように青・緑・ピンクの中でどこにおいてもかわらないのでindexが奇数の場合は-1引くなりして偶数にしておくと扱いやすい。
これを先生との差をとる生徒は必ず偶数にいるというのもわかる。

つぎに身長の差は以下のように累積和を考えればO(1)で求められる。

先生の身長をtとすると、index=0に挿入するときはA_0 + |H_0 - t| + B_0で身長差の合計が求められる。

これで計算量はO(MlogN + NlogN)となったので、あとは上記をまとめてコードにすればよい。

コード

実装時のTips

  • 累積和の配列のindex計算が非常にややこしいので図にかいて数ケース考えてから実装するといい
#include <bits/stdc++.h>

#include <atcoder/all>

using namespace std;
using namespace atcoder;
using ll = long long;
using ld = long double;
using uint = unsigned int;
using ull = unsigned long long;
const int MOD = 1e9 + 7;

int main() {
  ll n, m;
  cin >> n >> m;
  vector<ll> h(n);
  vector<ll> w(m);
  for (int i = 0; i < n; i++) {
    cin >> h[i];
  }
  sort(h.begin(), h.end());
  for (int i = 0; i < m; i++) {
    cin >> w[i];
  }
  vector<int> a((n + 1) / 2);  // a[0] = 0
  vector<int> b((n + 1) / 2);  // b[(n+1)/2 -1] = 0

  for (int i = 0; i + 1 < n; i += 2) {  // 奇数なので i + 1
    int ai = i / 2 + 1;
    a[ai] = h[i + 1] - h[i] + a[ai - 1];
  }
  for (int i = n - 2; i > 0; i -= 2) {
    int bi = (i - 1) / 2;
    b[bi] = h[i + 1] - h[i] + b[bi + 1];
  }

  ll r = LONG_LONG_MAX;
  for (int i = 0; i < m; i++) {
    ll t = w[i];
    int ti = lower_bound(h.begin(), h.end(), t) - h.begin();
    // 先生と生徒の差を出しやすいようにインデックスを偶数にする
    ti -= ti % 2;
    ll v = a[ti / 2] + abs(h[ti] - t) + b[ti / 2];
    r = min(r, v);
  }
  cout << r << endl;
}

参考