作者:杉原厚吉
共立出版

背景

プログラミングテストの練習問題として、
優先度付きキューの一種であるヒープを利用した
Kth Largest Element in a Stream問題を解いたので、
その解法をメモしておく。

問題

入力として正の整数値がいくつか入ったベクトルがあるとする。
そのベクトルに別の整数を追加した際に、その中からK番目に
大きい要素を取り出すアルゴリズムを実装する。

今回は、Kを3、入力するベクトルを[4, 5, 8, 2]とする。そして、このベクトルに3, 5, 10, 9, 4といった5つの数値を
入れた際のK番目に大きな要素を求める。

解法

ここでは、ヒープを利用したアルゴリズムを実装する。
ヒープというデータ構造についてはこちらの記事が詳しい。
medium.com

そして、今回実装するアルゴリズムは下記のようになる。

インスタンス生成時に、入力ベクトルのmin-heapを生成する。この時、ヒープの要素数がKとなるように、要素を小さいものから順に削除する(pop)
ベクトルに数値を入れたら、その時点での要素数がKを超えていないかチェックする。超えていたら、popによって要素数がKになるまで小さい順に削除する。その後、残った要素の中で一番小さいもの取り出せば、それが求めるＫ番目に大きい要素となる。

C++におけるmin-heapの実装方法はこちらを参照。

cpprefjp.github.io

www.geeksforgeeks.org

今回実装した上記のアルゴリズムのC++コードは下記の通り。

#include <iostream>
#include <queue>

using namespace std;

class KthLargest {
  private:
    int k;
    priority_queue <int, vector<int>, greater<int>> min_heap;
  public:
    KthLargest(int k, vector<int>& nums) {
      this->k = k;

      for (int num : nums) {
        this->min_heap.push(num);
      }

      while (this->min_heap.size() > this->k) {
        this->min_heap.pop();
      }
    }

    int add(int val) {
      this->min_heap.push(val);

      if (this->min_heap.size() > this->k) {
        this->min_heap.pop();
      }

      return this->min_heap.top();
    }
};

int main(void) {
  vector<int> input = {4, 5, 8, 2};
  int k = 3;

  KthLargest kl = KthLargest(k, input);

  int result_test1 = kl.add(3);
  cout << "Add 3: Kth largest num is " << result_test1 << endl;

  int result_test2 = kl.add(5);
  cout << "Add 5: Kth largest num is " << result_test2 << endl;

  int result_test3 = kl.add(10);
  cout << "Add 10: Kth largest num is " << result_test3 << endl;

  int result_test4 = kl.add(9);
  cout << "Add 9: Kth largest num is " << result_test4 << endl;

  int result_test5 = kl.add(4);
  cout << "Add 4: Kth largest num is " << result_test5 << endl;
    
  return 0;
}