Do a better RandomizedSet

leetcode/stuff.py

@@ -1,4 +1,5 @@
 import copy
+import random
 import typing
 
 import stdlib
@@ -393,3 +394,54 @@ def h_index(citations: list[int]) -> int:
             break
 
     return last_qualified or 0
+
+
+class SlowRandomizedSet:
+    def __init__(self):
+        self._i: set[int] = set()
+
+    def insert(self, val: int) -> bool:
+        ok = val not in self._i
+        self._i.add(val)
+        return ok
+
+    def remove(self, val: int) -> bool:
+        if val in self._i:
+            self._i.remove(val)
+            return True
+        return False
+
+    def getRandom(self) -> int:
+        return random.choice(list(self._i))
+
+
+class RandomizedSet:
+    def __init__(self):
+        self._l: list[int] = []
+        self._m: dict[int, int] = {}
+
+    def insert(self, val: int) -> bool:
+        if val in self._m:
+            return False
+
+        self._m[val] = len(self._l)
+        self._l.append(val)
+
+        return True
+
+    def remove(self, val: int) -> bool:
+        if val not in self._m:
+            return False
+
+        val_loc = self._m[val]
+        last_val = self._l[-1]
+        self._l[val_loc] = last_val
+        self._m[last_val] = val_loc
+
+        self._l.pop()
+        self._m.pop(val)
+
+        return True
+
+    def getRandom(self) -> int:
+        return random.choice(self._l)

leetcode/test_stuff.py

@@ -331,3 +331,60 @@ def test_count_min_jumps_from_board(board: list[int], expected: int):
 )
 def test_h_index(citations: list[int], expected: int):
     assert stuff.h_index(citations) == expected
+
+
+@pytest.mark.parametrize(
+    ("cls",),
+    [
+        (stuff.SlowRandomizedSet,),
+        (stuff.RandomizedSet,),
+    ],
+)
+def test_randomized_set(cls: type[stuff.RandomizedSet] | type[stuff.SlowRandomizedSet]):
+    inst = cls()
+
+    assert inst.insert(1) is True
+    assert inst.remove(2) is False
+    assert inst.insert(2) is True
+    assert inst.getRandom() in (1, 2)
+    assert inst.remove(1) is True
+    assert inst.insert(2) is False
+    assert inst.getRandom() == 2
+
+    inst = cls()
+
+    assert inst.insert(1) is True
+    assert inst.insert(10) is True
+    assert inst.insert(20) is True
+    assert inst.insert(30) is True
+
+    seen: set[int] = set()
+
+    for _ in range(10_000):
+        seen.add(inst.getRandom())
+
+    assert seen == {1, 10, 20, 30}
+
+    # ["remove","remove","insert","getRandom","remove","insert"]
+    # [[0],[0],[0],[],[0],[0]]
+
+    inst = cls()
+
+    assert inst.remove(0) is False
+    assert inst.remove(0) is False
+    assert inst.insert(0) is True
+    assert inst.getRandom() == 0
+    assert inst.remove(0) is True
+    assert inst.insert(0) is True
+
+    # ["RandomizedSet","insert","insert","remove","insert","remove","getRandom"]
+    # [[],[0],[1],[0],[2],[1],[]]
+
+    inst = cls()
+
+    assert inst.insert(0) is True
+    assert inst.insert(1) is True
+    assert inst.remove(0) is True
+    assert inst.insert(2) is True
+    assert inst.remove(1) is True
+    assert inst.getRandom() == 2