LeetCode 703. Kth Largest Element in a Stream

プログラミング
Page content

どうも、たくチャレ(@takuchalle)です。

設問

最初にkと配列numsが与えられ、add関数の返り値としてk番目に大きい値を返す問題。

解

優先順位付きキューpriority_queueを使えば簡単に解くことができます。

優先順位付きキューを使い、昇順に並ぶようにし、キューの要素数をk個に保つことでtopを呼んだ時に常にk番目に大きい値を取得できるようにしています。

コンストラクタの計算量はO(NlogN)、addメソッドの計算量はO(logN)で、全体の使用メモリはO(N)です。

class KthLargest {
   private:
    priority_queue<int, vector<int>, greater<int>> _q;
    int _k;

   public:
    KthLargest(int k, vector<int>& nums) {
        _k = k;

        for (auto v : nums) {
            _q.push(v);
            if (_q.size() > _k) {
                _q.pop(); // 一番小さい値を除く
            }
        }
    }

    int add(int val) {
        _q.push(val);

        if (_q.size() > _k) {
            _q.pop();
        }

        return _q.top();
    }
};