ARC101 D - Median of Medians

解法

$a[i]$が中央値となる区間の数を求めて$m=\text{floor}(n*(n+1)/4)+1$番目の数を求める方針で考えてたけど全くできなかったので終了後解説を見た。
解法としては二分探索をする。中央値の中央値がx以上か？のクエリについて考えると単調性がある。このクエリを高速に解く方法について考えていく。

中央値の中央値が$x$以上か？
区間$[l,r]$で中央値が$x$未満の区間が$m$個未満か？
区間$[l,r]$で($x$未満の個数)>($x$以上の個数)の区間が$m$個未満か？
$x$未満に1、$x$以上に-1を割り振る、区間$[l,r]$の和が0より大きくなる区間が$m$個未満か？
$S[i]=($区間$[0,i]$の和$)$とすると$S[r]-S[l-1]>0$を満たす組の数が$m$個未満か？

最終的に数列$S$の$S[r]>S[l-1]$を満たす組を数える問題に帰着できる。数列を前から見ていき$S[i]$以下の個数をBITを用いて保持しておくことでこのクエリは$O(N\log N)$で解ける。まとめると$O(N\log N\log A)$となる。
公式の解説と本質は同じだけどクエリの置き方とかが違うので注意。境界条件が未満、以上などややこしいので実装には要注意。

学び

$K$番目の値が何かの問題で「$K$番目の値が$X$以上か？」みたいに二分探索に落とすのは典型
中央値を考えるときに$X$以上か$X$未満かの2種類に落とす

ソースコード

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103


#include <bits/stdc++.h>

using namespace std;
using ll = long long;
#define int ll
using PII = pair<int, int>;
template <typename T> using V = vector<T>;
template <typename T> using VV = vector<V<T>>;
template <typename T> using VVV = vector<VV<T>>;

#define FOR(i, a, n) for (ll i = (ll)a; i < (ll)n; ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(x) x.begin(), x.end()
#define PB push_back

const ll INF = (1LL<<60);
const int MOD = 1000000007;

template <typename T> T &chmin(T &a, const T &b) { return a = min(a, b); }
template <typename T> T &chmax(T &a, const T &b) { return a = max(a, b); }
template <typename T> bool IN(T a, T b, T x) { return a<=x&&x<b; }
template<typename T> T ceil(T a, T b) { return a/b + !!(a%b); }
template<class S,class T>
ostream &operator <<(ostream& out,const pair<S,T>& a){
  out<<'('<<a.first<<','<<a.second<<')';
  return out;
}
template<class T>
ostream &operator <<(ostream& out,const vector<T>& a){
  out<<'[';
  REP(i, a.size()) {out<<a[i];if(i!=a.size()-1)out<<',';}
  out<<']';
  return out;
}

int dx[] = {0, 1, 0, -1}, dy[] = {1, 0, -1, 0};

// Binary Indexed Tree
// 0-indexed
template <typename T>
class BIT {
private:
  // データ
  vector<T> bit;
  // 単位元, 要素数
  int neutral = 0;
  // 更新クエリ, 区間クエリ
  function<T(T,T)> f = [](const T l, const T r) -> T { return l+r; },
                   g = [](const T l, const T r) -> T { return l+r; };
public:
  // 初期化
  BIT(int n_ = 1e5) { init(n_); }
  void init(int n_ = 1e5) { bit.assign(n_+1, neutral); }
  // iに対する点更新クエリ
  void update(int i, T w) {
    for(int x = i+1; x < bit.size(); x += x&-x) bit[x] = f(bit[x], w);
  }
  // [0,i]に対する区間クエリ
  T query(int i) {
    T ret = neutral;
    for(int x = i+1; x > 0; x -= x & -x) ret = g(ret, bit[x]);
    return ret;
  }
};


signed main(void)
{
  cin.tie(0);
  ios::sync_with_stdio(false);

  int n;
  cin >> n;
  V<int> a(n);
  REP(i, n) cin >> a[i];

  auto check = [&](int mid) -> bool {
    int now = 1e5;
    int ret = 0;
    BIT<int> bit(2e5+10);
    bit.update(now, 1);
    REP(i, n) {
      if(a[i] < mid) now++;
      else now--;
      bit.update(now, 1);
      ret += bit.query(now-1);
    }
    int m = n*(n+1)/4 + 1;
    return ret < m;
  };

  int lb = 0, ub = 1e9+10;
  while(ub-lb > 1) {
    int mid = (lb+ub)/2;
    if(check(mid)) lb = mid;
    else ub = mid;
  }

  cout << lb << endl;

  return 0;
}