PythonでInsert Delete GetRandom O(1)を実装する方法｜平均O(1)で動作するデータ構造

本記事では、平均O(1)時間で以下の3つの操作をすべてサポートできるデータ構造を、Pythonで実装する方法を解説します。

• insert(val) — 要素 val がまだセットに存在しない場合に挿入します。
• remove(val) — 要素 val がセットに存在する場合に削除します。
• getRandom() — 現在の要素集合からランダムに1つの要素を返します。どの要素も等しい確率で選ばれる必要があります。

アルゴリズムの考え方

ポイントは、辞書（ハッシュマップ）とリスト（動的配列）を組み合わせることです。辞書によって要素の有無をO(1)で判定でき、リストの末尾操作により削除もO(1)で行えます。具体的な手順は以下の通りです。

• 初期化： 要素の有無を記録する辞書（present）と、要素を格納する配列（elements）を用意します。
• insert(val)： valが辞書に存在しない、または present[val] == 0 の場合、elementsの末尾にvalを追加し、present[val] = 1 を設定して True を返します。それ以外の場合は False を返します。
• remove(val)： valが辞書に存在しない、または present[val] == 0 の場合は False を返します。存在する場合は、まず present[val] = 0 とし、elements内でのvalのインデックスを取得します。そのインデックスが末尾でなければ、末尾の要素と入れ替えます。最後にelementsの末尾要素を削除して True を返します。
• getRandom()： elements配列からランダムに1つの値を選んで返します。

削除時に「対象要素を末尾の要素と入れ替えてからpopする」のが重要なテクニックです。リストの中間要素の削除はO(n)かかりますが、末尾の削除であればO(1)で完了するためです。

Pythonでの実装例

それでは、実際のコードを見てみましょう。

import random
class RandomizedSet(object):
   def __init__(self):
      self.present = {}
      self.elements = []
   def insert(self, val):
      if val not in self.present or self.present[val] == 0:
         self.elements.append(val)
         self.present[val] = 1
         return True
      return False
   def remove(self, val):
      if val not in self.present or self.present[val] == 0:
         return False
      self.present[val] = 0
      index = self.elements.index(val)
      if index != len(self.elements)-1:
         temp = self.elements[-1]
         self.elements[-1] = val
         self.elements[index] = temp
      self.elements.pop()
      return True
   def getRandom(self):
      return random.choice(self.elements)
ob = RandomizedSet()
print(ob.insert(1))
print(ob.remove(2))
print(ob.insert(2))
print(ob.getRandom())
print(ob.remove(1))
print(ob.insert(2))
print(ob.getRandom())

入力

クラスを初期化した後、insert()、remove()、getRandom()関数を順に呼び出します。詳細は上記の実装例をご覧ください。

出力

True
False
True
2
True
False
2

このように、辞書とリストを併用することで、挿入・削除・ランダム取得のすべての操作を平均O(1)時間で実現できます。要素の重複チェックには辞書、ランダムアクセスにはリスト、という役割分担がこのデータ構造の鍵となります。