LoginSignup
105
87

More than 5 years have passed since last update.

@drken(けんちょん (Otsuki))in株式会社NTTデータ数理システム

k 番目の値を高速に取り出せるデータ構造のまとめ - BIT上二分探索や平衡二分探索木など

Last updated at Posted at 2018-02-13

はじめに

ビックデータという言葉が流行ってから数年の歳月が経ち、耳にする機会が大分減ってしまいましたが、現在も大規模データを高速に処理していくことは様々な現場で重要な課題となっています。

今回はその一端として、以下の操作を高速に実現するデータ構造について特集します:

  • 集合に要素 $v$ を挿入する
  • 集合から要素 $v$ を削除する
  • 集合に含まれる $k$ 番目に小さな値を取得する

「挿入」や「削除」を行えるデータ構造というとリストが真っ先に思い浮かびますが、「$k$ 番目に小さな値を求めよ」と言われると一気に難易度が上がります。ついつい平衡二分探索木を使いたくなるのですが、BIT や priority_queue で対応できるケースも多いです。BIT や平衡二分探索木の詳細についてはとてもよい資料が多数あるので以下に示します。

BIT

BIT については hos さんの資料

に極めてよくまとまっています。BIT の基本から二次元 BIT、区間加算対応 BIT、BIT 上の二分探索で $k$ 番目の値を取得するところまでカバーしています。

平衡二分探索木

平衡二分探索木についての資料としては、

がとても読みやすいです。また自前で実装を用意するのが大変でも G++ 拡張の Policy Based Data Structure の一つに tree と呼ばれるものがあるようです。$k$ 番目に小さな値を取得することもできるようです。使い方は hogloid さんの以下の記事にまとまっています。

　

整理

要素の値 v が 1 〜 D に限られているとき (D = 100000 程度)

BIT を用いることで実現できます。BIT の各 index における値として、その index が何個 S に含まれているかをもちます。

  • $v$ を挿入: bit.add(v, 1)
  • $v$ を削除: bit.add(v, -1)
  • $k$ 番目を求める: bit.sum(v) >= k となる最小の v を二分探索によって求める

という感じです。$k$ 番目を求める二分探索法は、ナイーブな実装では $O((\log D)^2)$ ですが、BIT上の二分探索を用いると $O(\log D)$ になります (hos さんの資料参照)。

  • 挿入: $O(\log D)$
  • 削除: $O(\log D)$
  • $k$ 番目: $O(\log D)$

なお他の方法として平方分割による解法も考えられます。区間 $[1, D]$ を $\sqrt{D}$ 個の区間 $[1, \sqrt{D}], [\sqrt{D}+1, 2\sqrt{D}], [2\sqrt{D}+1, 3\sqrt{D}], ..., [(\sqrt{D}-1)\sqrt{D}+1, D]$ に分割して処理します。そのようにすると「挿入」「削除」「$k$ 番目取得」はいずれも $O(\sqrt{D})$ で実施できます。

要素 v のとりうる値をすべて先読みできるとき

座標圧縮からの BIT で対応できます。集合に追加する要素が (a[1], a[2], ..., a[N]) の N 個だと予めわかっていれば、それを小さい順に並べて 1, 2, ..., N と番号付けしてあげれば、結局上の場合と同様に考えられます。

  • 挿入: $O(\log N)$
  • 削除: $O(\log N)$
  • $k$ 番目: $O(\log N)$

k が固定のとき

2 つの priority_queue で対応できます。あるいは $k$ が固定でなくても median を取り出すクエリも同様に対応できます。

  • 挿入: $O(\log N)$
  • 削除: $O(\log N)$
  • $k$ 番目: $O(1)$

どうにも工夫困難なとき

平衡二分探索木を用いましょう。

  • 挿入: $O(\log N)$
  • 削除: $O(\log N)$
  • $k$ 番目: $O(\log N)$

問題例

ARC 033 C - データ構造

数の集合 S に対する以下の Q 個のクエリを処理してください。i 番目のクエリは以下のいずれかです。

  • タイプ 1： S に数 X[i] を追加する。
  • タイプ 2： S に含まれる数のうち X[i] 番目に小さい数を答え、その数を S から削除する。

・1 <= Q <= 200000
・1 <= X[i] <= 200000

上のうち「要素の値 v が 1 〜 D に限られているとき」に該当します。
ソースコードにすると次のようになります。

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

template <class Abel> struct BIT {
    vector<Abel> dat;
    Abel UNITY_SUM = 0;                     // to be set

    /* [1, n] */
    BIT(int n) { init(n); }
    void init(int n) {
        dat.resize(n + 1);
        for (int i = 0; i < (int)dat.size(); ++i) dat[i] = UNITY_SUM;
    }

    /* a is 1-indexed */
    inline void add(int a, Abel x) {
        for (int i = a; i < (int)dat.size(); i += i & -i)
            dat[i] = dat[i] + x;
    }

    /* [1, a], a is 1-indexed */
    inline Abel sum(int a) {
        Abel res = UNITY_SUM;
        for (int i = a; i > 0; i -= i & -i)
            res = res + dat[i];
        return res;
    }

    /* [a, b), a and b are 1-indexed */
    inline Abel sum(int a, int b) {
        return sum(b - 1) - sum(a - 1);
    }

    /* k-th number (k is 0-indexed) */
    int get(long long k) {
        ++k;
        int res = 0;
        int N = 1; while (N < (int)dat.size()) N *= 2;
        for (int i = N / 2; i > 0; i /= 2) {
            if (res + i < (int)dat.size() && dat[res + i] < k) {
                k = k - dat[res + i];
                res = res + i;
            }
        }
        return res + 1;
    }
};

int main() {
    BIT<int> bit(200000);
    int Q;
    cin >> Q;
    for (int i = 0; i < Q; ++i) {
        int T, X;
        cin >> T >> X;
        if (T == 1) bit.add(X, 1); // insert X
        else {
            int val = bit.get(X - 1); // X is 1-indexed
            cout << val << endl;
            bit.add(val, -1); // delete val
        }
    }
}

　

yukicoder No.649 ここでちょっとQK！

K は固定で与えられる。数の集合 S に対する以下の Q 個のクエリを処理してください。i 番目のクエリは以下のいずれかです。

  • タイプ 1： S に数 v[i] を追加する。
  • タイプ 2： S に含まれる数のうち K 番目に小さい数を答え、その数を S から削除する。

・1 <= Q <= 200000
・1 <= K <= 200000
・0 <= v[i] <= 10^18

今度は v[i] の値が 10^18 までありうるので、さきほどと同じ手は使えません。この問題に対しては以下のような解法が可能です！

  • 要素 v のとりうる値をすべて先読みできるとき: 座標圧縮して BIT
  • $k$ が固定のとき: 2 つの priority_queue
  • 平衡二分探索木で殴る

1 つずつ実装していきたいと思います。

　

座標圧縮して BIT

ほとんど先ほどの ARC の問題と一緒なのですが、ありうる v[i] の値をすべて先読みして座標圧縮をしています。

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

template <class Abel> struct BIT {
    vector<Abel> dat;
    Abel UNITY_SUM = 0;                     // to be set

    /* [1, n] */
    BIT(int n) { init(n); }
    void init(int n) {
        dat.resize(n + 1);
        for (int i = 0; i < (int)dat.size(); ++i) dat[i] = UNITY_SUM;
    }

    /* a is 1-indexed */
    inline void add(int a, Abel x) {
        for (int i = a; i < (int)dat.size(); i += i & -i)
            dat[i] = dat[i] + x;
    }

    /* [1, a], a is 1-indexed */
    inline Abel sum(int a) {
        Abel res = UNITY_SUM;
        for (int i = a; i > 0; i -= i & -i)
            res = res + dat[i];
        return res;
    }

    /* [a, b), a and b are 1-indexed */
    inline Abel sum(int a, int b) {
        return sum(b - 1) - sum(a - 1);
    }

    /* k-th number (k is 0-indexed) */
    int get(long long k) {
        ++k;
        int res = 0;
        int N = 1; while (N < (int)dat.size()) N *= 2;
        for (int i = N / 2; i > 0; i /= 2) {
            if (res + i < (int)dat.size() && dat[res + i] < k) {
                k = k - dat[res + i];
                res = res + i;
            }
        }
        return res + 1;
    }
};


int main() {
    int max_Q = 210000;
    BIT<int> bit(max_Q);
    int Q, K;
    cin >> Q >> K;
    vector<int> T(Q);
    vector<long long> V(Q);

    // クエリ先読み
    vector<long long> queries;
    for (int i = 0; i < Q; ++i) {
        cin >> T[i];
        if (T[i] == 1) cin >> V[i], queries.push_back(V[i]);
    }

    // 座標圧縮
    sort(queries.begin(), queries.end());
    queries.erase(unique(queries.begin(), queries.end()), queries.end());

    // クエリに答える
    for (int i = 0; i < Q; ++i) {
        if (T[i] == 1) {
            int index = lower_bound(queries.begin(), queries.end(), V[i]) - queries.begin();
            ++index; // 1-index にする
            bit.add(index, 1); // insert index
        }
        else {
            int size = bit.sum(max_Q); // insert されている全要素数
            if (K > size) {
                cout << -1 << endl;
            }
            else {
                int val = bit.get(K - 1); // k-th number (original K is 1-indexed)
                cout << queries[val - 1] << endl; // val を 0-indexed にする
                bit.add(val, -1); // erase val
            }
        }
    }
}

2 つの priority_queue

この yukicoder の問題は、v[i] が 10^18 までありうるという点では ARC 030 C より厳しい問題ですが、$k$ が固定という点では優しいです。2 つの priority_queue を使うことで $k$ 番目の値を常に取り出せる状態を確保できます。

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

int main() {
  // small.top() は small の最大値 (K個になるようにする)
  priority_queue<long long> small;

  // large.top() は large の最小値
  priority_queue<long long, vector<long long>, greater<long long> > large;

  int Q, K;
  cin >> Q >> K;
  for (int i = 0; i < Q; ++i) {
    int T;
    cin >> T;
    if (T == 1) {
      long long V;
      cin >> V;

      // そもそも K 個に達していなかったら small へ
      if (small.size() < K) {
        small.push(V);
        continue;
      }

      // k 番目の値
      long long current_kth = small.top();

      // 現在の k 番目の値より小さかったら、それをlarge 側へ追い出して V を small に push
      if (V < current_kth) {
        small.pop();
        small.push(V);
        large.push(current_kth);
      }
      // そうでないときは単純に large 側に push
      else {
        large.push(V);
      }
    }
    else {
      // そもそも K 個に達していなかったらダメ
      if (small.size() < K) {
        cout << -1 << endl;
        continue;
      }

      // k 番目の値を small から pop して、large の top を small へ移す
      long long current_kth = small.top();
      cout << current_kth << endl;
      small.pop();
      if (!large.empty()) {
        long long next_kth = large.top();
        large.pop();
        small.push(next_kth);
      }
    }
  }
}

平衡二分探索木

伝家の宝刀、平衡二分探索木で殴ることもできます。「挿入」「削除」「$k$ 番目の値」に限らず、他にも様々なことを log オーダーでこなすことができます。ここでは、merge-split ベースの RBST (Randomized Binary Search Tree) を実装しています。

#include <iostream>
using namespace std;

template<class VAL> struct RBST {
    VAL SUM_UNITY = 0;                              // to be set

    unsigned int randInt() {
        static unsigned int tx = 123456789, ty = 362436069, tz = 521288629, tw = 88675123;
        unsigned int tt = (tx ^ (tx << 11));
        tx = ty; ty = tz; tz = tw;
        return (tw = (tw ^ (tw >> 19)) ^ (tt ^ (tt >> 8)));
    }

    struct NODE {
        NODE *left, *right;
        VAL val;                        // the value of the node
        int size;                       // the size of the subtree 
        VAL sum;                        // the value-sum of the subtree

        NODE() : val(SUM_UNITY), size(1), sum(SUM_UNITY) {
            left = right = NULL;
        }

        NODE(VAL v) : val(v), size(1), sum(v) {
            left = right = NULL;
        }

        /* additional update */
        inline void update() {

        }

        /* additional lazy-propagation */
        inline void push() {

            /* ex: reverse */
            /*
            if (this->rev) {
            swap(this->left, this->right);
            if (this->left) this->left->rev ^= true;
            if (this->right) this->right->rev ^= true;
            this->rev = false;
            }
            */
        }
    };


    ///////////////////////
    // root
    ///////////////////////

    NODE* root;
    RBST() : root(NULL) { }
    RBST(NODE* node) : root(node) { }


    ///////////////////////
    // basic operations
    ///////////////////////

    /* size */
    inline int size(NODE *node) {
        return !node ? 0 : node->size;
    }
    inline int size() {
        return this->size(this->root);
    }

    /* sum */
    inline VAL sum(NODE *node) {
        return !node ? SUM_UNITY : node->sum;
    }
    inline VAL sum() {
        return this->sum(this->root);
    }

    /* update, push */
    inline NODE* update(NODE *node) {
        node->size = size(node->left) + size(node->right) + 1;
        node->sum = sum(node->left) + sum(node->right) + node->val;
        node->update();
        return node;
    }

    inline void push(NODE *node) {
        if (!node) return;
        node->push();
    }

    /* lower_bound */
    inline int lowerBound(NODE *node, VAL val) {
        push(node);
        if (!node) return 0;
        if (val <= node->val) return lowerBound(node->left, val);
        else return size(node->left) + lowerBound(node->right, val) + 1;
    }
    inline int lowerBound(VAL val) {
        return this->lowerBound(this->root, val);
    }

    /* upper_bound */
    inline int upperBound(NODE *node, VAL val) {
        push(node);
        if (!node) return 0;
        if (val >= node->val) return size(node->left) + upperBound(node->right, val) + 1;
        else return upperBound(node->left, val);
    }
    inline int upperBound(VAL val) {
        return this->upperBound(this->root, val);
    }

    /* count */
    inline int count(VAL val) {
        return upperBound(val) - lowerBound(val);
    }

    /* get --- k: 0-index */
    inline VAL get(NODE *node, int k) {
        push(node);
        if (!node) return -1;
        if (k == size(node->left)) return node->val;
        if (k < size(node->left)) return get(node->left, k);
        else return get(node->right, k - size(node->left) - 1);
    }
    inline VAL get(int k) {
        return get(this->root, k);
    }


    ///////////////////////
    // merge-split
    ///////////////////////

    NODE* merge(NODE *left, NODE *right) {
        push(left);
        push(right);
        if (!left || !right) {
            if (left) return left;
            else return right;
        }
        if (randInt() % (left->size + right->size) < left->size) {
            left->right = merge(left->right, right);
            return update(left);
        }
        else {
            right->left = merge(left, right->left);
            return update(right);
        }
    }
    void merge(RBST add) {
        this->root = this->merge(this->root, add.root);
    }
    pair<NODE*, NODE*> split(NODE* node, int k) { // [0, k), [k, n)
        push(node);
        if (!node) return make_pair(node, node);
        if (k <= size(node->left)) {
            pair<NODE*, NODE*> sub = split(node->left, k);
            node->left = sub.second;
            return make_pair(sub.first, update(node));
        }
        else {
            pair<NODE*, NODE*> sub = split(node->right, k - size(node->left) - 1);
            node->right = sub.first;
            return make_pair(update(node), sub.second);
        }
    }
    RBST split(int k) {
        pair<NODE*, NODE*> sub = split(this->root, k);
        this->root = sub.first;
        return RBST(sub.second);
    }


    ///////////////////////
    // insert-erase
    ///////////////////////

    void insert(const VAL val) {
        pair<NODE*, NODE*> sub = this->split(this->root, this->lowerBound(val));
        this->root = this->merge(this->merge(sub.first, new NODE(val)), sub.second);
    }

    void erase(const VAL val) {
        if (!this->count(val)) return;
        pair<NODE*, NODE*> sub = this->split(this->root, this->lowerBound(val));
        this->root = this->merge(sub.first, this->split(sub.second, 1).second);
    }


    ///////////////////////
    // debug
    ///////////////////////

    void print(NODE *node) {
        if (!node) return;
        push(node);
        print(node->left);
        cout << node->val << " ";
        print(node->right);
    }
    void print() {
        cout << "{";
        print(this->root);
        cout << "}" << endl;
    }
};

int main() {
    RBST<long long> S;
    int Q, K;
    cin >> Q >> K;
    for (int i = 0; i < Q; ++i) {
        int T;
        cin >> T;
        if (T == 1) {
            long long val;
            cin >> val;
            S.insert(val);
        }
        else {
            if (S.size() < K) {
                cout << -1 << endl;
                continue;
            }
            long long res = S.get(K - 1);
            cout << res << endl;
            S.erase(res);
        }
        //S.print();
    }
}
105
87
2

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
105
87