From c8a4928ee8a3e328f8502ea063998c2cdcef880c Mon Sep 17 00:00:00 2001
From: Dan Buch <dan@meatballhat.com>
Date: Wed, 25 Oct 2023 20:15:43 -0400
Subject: [PATCH] a rushed RandomizedSet impl

---
 leetcode/stuff.py      | 20 ++++++++++++++++++++
 leetcode/test_stuff.py | 26 ++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/leetcode/stuff.py b/leetcode/stuff.py
index 580b18e..d98704a 100644
--- a/leetcode/stuff.py
+++ b/leetcode/stuff.py
@@ -1,4 +1,5 @@
 import copy
+import random
 import typing
 
 import stdlib
@@ -393,3 +394,22 @@ def h_index(citations: list[int]) -> int:
             break
 
     return last_qualified or 0
+
+
+class RandomizedSet:
+    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 next(iter(sorted(list(self._i), key=lambda _: random.random())))
diff --git a/leetcode/test_stuff.py b/leetcode/test_stuff.py
index 3ba42a3..3730f43 100644
--- a/leetcode/test_stuff.py
+++ b/leetcode/test_stuff.py
@@ -331,3 +331,29 @@ 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
+
+
+def test_randomized_set():
+    inst = stuff.RandomizedSet()
+
+    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 = stuff.RandomizedSet()
+
+    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(100_000):
+        seen.add(inst.getRandom())
+
+    assert seen == {1, 10, 20, 30}