区間のk-th smallestが求められる謎データ構造 - 競技プログラミング+αなブログ

この記事は
Competitive Programming Advent Calendar Div2013(http://partake.in/events/3a3bb090-1390-4b2a-b38b-4273bea4cc83)
の16日目の記事です。

データ構造の話です。

地区予選の過去問を解いているときに、以下のような問題を見つけました。

【問題】
要素数Nの数列a(a[0],...,a[N-1])が与えられる。
以下の２種類からなるクエリがQ個与えられるので、それぞれ処理せよ。
・Query1 l r k
　a[l],...,a[r-1]のk番目に小さい数は何か　（1<=k<=r-l）
・Query2 l r x
　a[l],...,a[r-1]をソートしたとき、xは何番目か　（x∈{a[l], ..., a[r]}）
【制約】
N<=10^5
Q<=10^5
0<=l

僕は解き方がわからなかったのですが、中国のサイトを漁っていたらこの問題を解くデータ構造を見つけました。
ということでそのデータ構造を紹介しようと思います。
そのデータ構造の名前はよくわからなかったので，今回は謎データ構造と呼ぶことにします．

この記事では，Query1について詳しく説明します．
Query2についてもQuery1と同じような考えを用いて求めることができます．
詳しくは最後に載せたソースコードをご覧ください．

例

例えば、aを

a[0]	a[1]	a[2]	a[3]	a[4]	a[5]	a[6]	a[7]
6	12	5	17	10	2	7	3

とします。
いま，
Query1 2 7 3
が与えられたとします．
まず，[2,7)以外の部分はいらないので，わかりやすくするために値を消しておきます．

a[0]	a[1]	a[2]	a[3]	a[4]	a[5]	a[6]	a[7]
		5	17	10	2	7

続いて，[2,7)をソートします．

a[0]	a[1]	a'[2]	a'[3]	a'[4]	a'[5]	a'[6]	a[7]
		2	5	7	10	17

ソートされた[2,7)の中で，3番目にある7が，このクエリに対する答えとなります．
ソートを全てのクエリに対して毎回行うと，O(QN log N)となり間に合いません．

アイデア

謎データ構造のアイデアは，求める値がどこにあるかを，二分探索のように狭めていくことです．
今，
[L,R)の中で[l,r)に注目しており，この中でk番目の値を探している
とします．
先程の例では，
[0,8)の中で[2,7)に注目しており，この中で3番目の値を探している
となります．
謎データ構造では，[L,R)をどんどん半分にしていくことで，探索の範囲を狭めていき，1クエリに対しO(log N)で答えを出します．
ここで必要となるのが，

M=(L+R)/2としたとき，[L,M)を見るべきか，[M,R)を見るべきか
次のl,r,kはどうなるか

という処理です．
1.を行うため，[l,r)内でk番目の値が左半分にあるか右半分にあるかを高速に判断する必要があります．
この処理を元の数列のまま行うのは無理なので，次の段階に移るとき，[L,R)のうち，値の大きさが下半分になるものを[L,M)，上半分になるものを[M,R)に動かします．
ただし，[L,M),[M,R)内ではもとの順番は保つようにします．

例示をわかりやすくするため，ここで，
b[i][j] = 段階iでの列におけるj番目の値
とします．
先程の例を用いると，段階0では

b[0][0]	b[0][1]	b[0][2]	b[0][3]	b[0][4]	b[0][5]	b[0][6]	b[0][7]
6	12	5	17	10	2	7	3

となり，段階1では

b[1][0]	b[1][1]	b[1][2]	b[1][3]	b[1][4]	b[1][5]	b[1][6]	b[1][7]
6	5	2	3	12	17	10	7

となります．
全体を図示したものが以下です．

最後のb[3]は元の列をソートしたものになります．
ここでさらに，
left[i][j] = 段階iにおけるj番目の値が所属するブロックにおいて，次の段階で左にいったものの数のjまでの累積和
を導入します．
leftを図に追加すると以下になります．

ここで，
[0,8)の中で[2,7)に注目しており，この中で3番目の値を探している
とき，次に[0,4)を見るべきか，[4,8)を見るべきかという問題に戻ります．
left[0][7-1]-left[0][2-1]=2より，[2,7)から左のブロックに進んだものは2個あることがわかります．
今3番目の値を探しているので，次に見るべきは左のブロックではなく，右のブロックであることがわかります．
「次のl,r,kはどうなるか」という問題も，leftを利用することで求まります．
この場合は，l=5,r=8,k=1となります．

以下[L,R)がだた１つの要素に対応するようになるまで繰り返すことで，Query1の答えを得ることができます．

bとleftの計算はO(N log N)で可能なので，O(N log N)の前処理をすることで，各Query1に対してO(log N)で求められるデータ構造を作ることができました．

ソースコード

今回の記事で扱った問題は，a[i]が互いに異なることを想定していましたが，その条件がなくても問題ありません．
以下の実装では，（一部自信がないですが）値が同じa[i]があっても動くようになっているはずです．

#include <iostream>
#include <algorithm>
using namespace std;

const int N = 10000;

struct KthSmallest {
  struct Seg {
    int l, r, mid;
    void set(int _l, int _r) {
      l = _l; r = _r;
      mid = l+r>>1;
    }
  } seg[N<<2];
  int b[25][N], left[25][N], sorted[N];
  void init(int *a, int n) {
    for (int i=0; i<n; ++i) b[0][i] = sorted[i] = a[i];
    sort(sorted, sorted+n);
    build(0,n,0,1);
  }
  void build(int l, int r, int d, int idx) {
    seg[idx].set(l,r);
    if (l+1 == r) return;
    int mid = seg[idx].mid;
    int lsame = mid - l;
    for (int i=l; i<r; ++i)
      if (b[d][i] < sorted[mid])
        lsame--;
    int lpos = l, rpos = mid, same = 0;

    for (int i=l; i<r; ++i) {
      left[d][i] = (i!=l ? left[d][i-1] : 0);
      if (b[d][i] < sorted[mid]) {
        left[d][i]++;
        b[d+1][lpos++] = b[d][i];
      } else if (b[d][i] > sorted[mid]) {
        b[d+1][rpos++] = b[d][i];
      } else {
        if (same < lsame) {
          same++;
          left[d][i]++;
          b[d+1][lpos++] = b[d][i];
        } else {
          b[d+1][rpos++] = b[d][i];
        }
      }
    }
    build(l,mid,d+1,idx<<1);
    build(mid,r,d+1,idx<<1|1);
  }
  /*
    [l,r)をソートしたときk番目に来る値は何か
   */
  int kth(int l, int r, int k, int d=0, int idx=1) { // k : 1-origin!!!
    if (l+1 == r) return b[d][l];
    int ltl = (l!=seg[idx].l ? left[d][l-1] : 0);
    int tl = left[d][r-1] - ltl;

    if (tl >= k) {
      int newl = seg[idx].l + ltl;
      int newr = seg[idx].l + ltl + tl;
      return kth(newl,newr,k,d+1,idx<<1);
    } else {
      int mid = seg[idx].mid;
      int tr = r - l - tl;
      int ltr = l - seg[idx].l - ltl;
      int newl = mid + ltr;
      int newr = mid + ltr + tr;
      return kth(newl,newr,k-tl,d+1,idx<<1|1);
    }
  }
  /*
    [l,r)をソートしたときxは何番目に来るか．
    xが2つ以上あるときは，最後のもののrankを返す．
    xがないときはx未満で最大なもののrankを返す．
    x未満がないときは0を返す．
  */ 
  int rank(int l, int r, int x, int d=0, int idx=1) {
    if (seg[idx].l+1 == seg[idx].r) {
      return l+1==r&&sorted[l]<=x;
    }
    int ltl = (l!=seg[idx].l ? left[d][l-1] : 0);
    int tl = left[d][r-1] - ltl;
    int mid = seg[idx].mid;
    if (x < sorted[mid]) {
      int newl = seg[idx].l + ltl;
      int newr = seg[idx].l + ltl + tl;
      return rank(newl,newr,x,d+1,idx<<1);
    } else {
      int tr = r - l - tl;
      int ltr = l - seg[idx].l - ltl;
      int newl = mid + ltr;
      int newr = mid + ltr + tr;
      return tl + rank(newl,newr,x,d+1,idx<<1|1);
    }
  }
  /*
    [l,r)にxは何個あるか
  */ 
  int freq(int l, int r, int x) {
    return rank(l,r,x)-rank(l,r,x-1);
  }
} kth;

int main() {
  int a[8] = {6,12,5,17,10,2,7,3};
  kth.init(a, 8);
  cout << kth.kth(2,7,3) << endl; // 7
  cout << kth.rank(2,7,7) << endl; // 3
}

verify:
（地区予選問題のネタバレ注意）https://icpcarchive.ecs.baylor.edu/index.php?option=com_onlinejudge&Itemid=8&category=393&page=show_problem&problem=3035
ただし，同じ値が含まれない場合のみ

まとめ

本記事では，区間のkth-smallestを求められる謎データ構造を紹介しました．
このデータ構造の名前を知っていたり，「これ，○○で出来るよ！」みたいなのがあったりしたら是非教えて下さい．