Leetcode 703. Kth Largest Element in a Stream

Problem

数字が次々と与えられた時に(stream) 、その中でk番目に大きい数を見つけるクラスを作成します。

Design a class to find the kth largest element in a stream. Note that it is the kth largest element in the sorted order, not the kth distinct element.

この問題でいう、"stream"は数値の連続的なシーケンスを指します。つまり、時間とともに新しい数値が次々と追加されるような状況を想定しています。ここで重要なのは、全てのデータが最初から利用可能なわけではないという点です。

InputとOutputの例は次になります。

Input
["KthLargest", "add", "add", "add", "add", "add"]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
Output
[null, 4, 5, 5, 8, 8]

Explanation
KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]);
kthLargest.add(3);   // return 4
kthLargest.add(5);   // return 5
kthLargest.add(10);  // return 5
kthLargest.add(9);   // return 8
kthLargest.add(4);   // return 8

Bad Approach

一度配列をソートした後、二分探索などを用いてソートされた配列に対して適切な位置に要素を挿入する方法も考えられます。しかし、この方法は効率が悪いという欠点があります。

配列をソートするには、一般的にO(NlogN)の時間がかかります（ここでNは配列の要素数）。そして、新たな要素が追加する際には、挿入位置を見つけるのに O(logN)、実際にその位置に要素を挿入する操作はO(N)の時間がかかります。挿入位置以降の全ての要素を1つずつ後ろにずらす必要があるからです。したがって、新しい要素の追加には合計でO(logN) + O(N) = O(N)の時間が必要となります。

これをトータルの計算量に適用すると、addメソッドが最大でM回呼び出されるとすると、全体の計算量はO(NlogN + M*N)となります。

問題の制約によるとMとNはともに最大で10^4なので、最悪の場合の計算量はO(10^4 * 10^4) = O(10^8)となります。

Key Idea

この問題の主な課題は、データがストリームとして常に追加されるときに、いつでもk番目に大きい要素を効率的に見つけることができるようなデータ構造を設計することです。

1つめのポイントとしては、ヒープを使うということです。ヒープを使うと効果的な理由は、ヒープがそれぞれの操作（要素の追加、最小（または最大）要素の取得と削除）を高速に行うことができるからです。特に最小ヒープを使うと、ヒープの最上部（根）には常に最小の要素が来るため、最小の要素の取得と削除がO(1)で可能となります。また、新たな要素の追加もログ時間（O(logN)）で行うことができます。

2つめのポイントとしては、k番目に大きい要素を見つけたいので、「サイズkの最小ヒープ」を使います。ヒープの中には常にk番目に大きい要素までの数字が入っていて、ヒープの根（最小の要素）がk番目に大きい数になります。
なぜ最大ヒープを使わないのでしょうか？最大ヒープを使うと、ヒープの根（最上部）には常に最大の要素が位置するため、k番目に大きい要素を取り出すには、ヒープから最大の要素をk-1回削除する必要があります。これは時間的に非効率で、addメソッドが呼び出されるたびにO(klogN)の時間がかかることになります。

Implementation

import heapq

class KthLargest:

    def __init__(self, k, nums):
        self.k = k
        self.minHeap = nums
        heapq.heapify(self.minHeap)

        while len(self.minHeap) > k:
            heapq.heappop(self.minHeap)

    def add(self, val):
        heapq.heappush(self.minHeap, val)
        # If heap size is more than k, remove smallest element
        if len(self.minHeap) > self.k:
            heapq.heappop(self.minHeap)
        # The kth largest value is at the root of the min-heap
        return self.minHeap[0]

ヒープの状態を書くと、下記のようになります。

heapq モジュール

Pythonのheapqモジュールはヒープキューアルゴリズムを提供します。ヒープキューは、各要素が特定のソート順（デフォルトでは小さい順）を維持するようなリストです。Pythonのheapqモジュールは、特に最小ヒープの操作に対する関数を提供しています。最小ヒープでは、最小の要素が常に根（インデックス0）に位置します。

以下に主な関数とその使い方を示します：

• heapq.heappush(heap, item): heapという名前のヒープにitemをプッシュ（追加）します。ヒープのプロパティ（親が子より小さい）はこの操作後も維持されます。

• heapq.heappop(heap): ヒープから最小の要素をポップ（削除）し、その要素を返します。ヒープのプロパティはこの操作後も維持されます。

• heapq.heapify(x): リストxをインプレースでヒープに変換します。

これらの関数は次のように使用できます：

import heapq

# create an empty heap
heap = []

# push values onto the heap
heapq.heappush(heap, 2)
heapq.heappush(heap, 3)
heapq.heappush(heap, 1)
heapq.heappush(heap, 4)

print(heap)  # output: [1, 2, 3, 4], it's a min heap, so the smallest element is at root

# pop the smallest value from the heap
smallest = heapq.heappop(heap)

print(smallest)  # output: 1
print(heap)  # output: [2, 4, 3], the next smallest element is at root now

# heapify a list
nums = [4, 1, 3, 5, 2, 6]
heapq.heapify(nums)

print(nums)  # output: [1, 2, 3, 5, 4, 6], the list is now a valid min heap

Pythonのheapqモジュールの各関数についての時間計算量と空間計算量は以下のようになります。

heapq.heapify:

• 時間計算量: O(n)。ここで、nはリストの要素数です。heapify操作は全ての要素に対して行われるため、要素数に比例した時間がかかります。
• 空間計算量: O(1)。heapify操作はリスト上でインプレース（追加のスペースを使わずに）行われるため、追加の空間は必要ありません。

heapq.heappop:

• 時間計算量: O(logn)。ここで、nはヒープの要素数です。ヒープから最小要素を削除した後に、ヒープの特性を維持するためには、ヒープの高さに比例した時間がかかります（ヒープの高さはlogn）。
• 空間計算量: O(1)。heappop操作はヒープ上でインプレース（追加のスペースを使わずに）行われるため、追加の空間は必要ありません。

heapq.heappush:

• 時間計算量: O(logn)。ここで、nはヒープの要素数です。新しい要素をヒープに追加した後に、ヒープの特性を維持するためには、ヒープの高さに比例した時間がかかります（ヒープの高さはlogn）。
• 空間計算量: O(1)。heappush操作はヒープ上でインプレース（追加のスペースを使わずに）行われるため、追加の空間は必要ありません。

Complexity Analysis

InitとAddそれぞれにおいて、計算量は下記の様になります。

	Time complexity	Space Complexity
Init	O(N) + (N-k)xO(log(N)) =>O(Nlog(N)	O(N)
Add	M x log(k)	O(k)

まとめると、Time ComplexityはO(Nxlog(N) + Mxlog(k)、Space Complexityは、O(N)となります。

Reference

https://www.youtube.com/watch?v=hOjcdrqMoQ8
https://medium.com/@yasufumy/data-structure-heap-ecfd0989e5be
https://www.youtube.com/watch?v=t0Cq6tVNRBA