公式

D - 中央値チェック / Median Check 解説 by kyopro_friends


この問題は Fenwick treeやセグメントツリーのようなデータ構造を用いて解くことができます。

この問題を解くためには、multiset のように要素の追加・削除が高速に行えることに加え、「\(x\) 以下の要素が何個?」のような質問に高速に答えることができるデータ構造が必要があります。

これは値の範囲が小さければ「\(i\) が何個あるか」を表す配列をセグメントツリーに乗せることで実現できます。値の範囲が大きな今回も、クエリ先読み+座標圧縮により同様に処理することができます。

実装例 (C++)

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

int main(){
  int q;
  cin >> q;
  vector<pair<char,int>>query(q);
  for(int i=0; i<q; i++) cin >> query[i].first >> query[i].second;
  
  set<int>s;
  for(int i=0; i<q; i++){
    s.insert(query[i].second);
  }
  vector<int>i2v(s.begin(), s.end());
  map<int,int>v2i;
  for(int i=0; i<i2v.size(); i++){
    v2i[i2v[i]] = i;
  }
  
  int ans = 0;
  int n = 0;
  atcoder::fenwick_tree<int>seg(i2v.size());
  for(int i=0; i<q; i++){
    if(query[i].first == '+'){
      seg.add(v2i[query[i].second], 1);
      n++;
    }else{
      int l = seg.sum(0, v2i[query[i].second]);
      int r = seg.sum(0, v2i[query[i].second] + 1);
      if(l < (n+1) / 2 && (n+1) / 2 <= r){
        ans++;
      }
      seg.add(v2i[query[i].second], -1);
      n--;
    }
  }
  cout << ans << endl;
}

実装例 (Python)

from atcoder.fenwicktree import FenwickTree

Q = int(input())
query = []
for _ in range(Q):
  c, x = input().split()
  query.append((c, int(x)))

i2v = sorted(set(x for _, x in query))
v2i = {v:i for i, v in enumerate(i2v)}

ans = 0
n = 0
seg = FenwickTree(len(i2v))
for c, x in query:
  if c == '+':
    seg.add(v2i[x], 1)
    n += 1
  else:
    l = seg.sum(0, v2i[x])
    r = seg.sum(0, v2i[x] + 1)
    if l < (n+1) // 2 <= r:
      ans += 1
    seg.add(v2i[x], -1)
    n -=1

print(ans)

なお、C++であればgcc拡張の ext/pb_ds/assoc_container.hpp にある __gnu_pbds::tree 、python であれば sortedcontainersSortedList により、要素の追加・削除・カウントを直接行えるデータ構造を使うこともできます。

実装例 (C++)

#include<bits/stdc++.h>
#include<ext/pb_ds/assoc_container.hpp>
using namespace std;
using namespace __gnu_pbds;

int main(){
  int q;
  cin >> q;
  
  int ans = 0;
  tree<
    pair<int,int>,
    null_type,
    less<pair<int,int>>,
    rb_tree_tag,
    tree_order_statistics_node_update
  > t;
  for(int i=0; i<q; i++){
    char c;
    int x;
    cin >> c >> x;
    if(c == '+'){
      t.insert({x, i});
    }else{
      auto med = t.find_by_order((t.size() - 1) / 2);
      if(med -> first == x){
        ans++;
      }
      auto it = t.lower_bound({x, -1});
      t.erase(it);
    }
  }
  cout << ans << endl;
}

実装例 (Python)

from sortedcontainers import SortedList

Q = int(input())

ans = 0
S = SortedList()
for _ in range(Q):
  C, X = input().split()
  X = int(X)
  if C == '+':
    S.add(X)
  else:
    if S[(len(S)-1) // 2] == X:
      ans += 1
    S.discard(X)

print(ans)

投稿日時:
最終更新: