Do a better RandomizedSet

leetcode/test_stuff.py

@@ -333,8 +333,15 @@ def test_h_index(citations: list[int], expected: int):
     assert stuff.h_index(citations) == expected
 
 
-def test_randomized_set():
-    inst = stuff.RandomizedSet()
+@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
@@ -344,7 +351,7 @@ def test_randomized_set():
     assert inst.insert(2) is False
     assert inst.getRandom() == 2
 
-    inst = stuff.RandomizedSet()
+    inst = cls()
 
     assert inst.insert(1) is True
     assert inst.insert(10) is True
@@ -353,7 +360,31 @@ def test_randomized_set():
 
     seen: set[int] = set()
 
-    for _ in range(100_000):
+    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