Improve stable random distribution

Author: Donaim

micall/tests/test_stable_random_distribution.py

@@ -22,8 +22,8 @@ def test_indices_in_range():
 def test_bounds_are_reachable():
     """Test that both min and max-1 can be generated."""
 
-    high = 999
-    rng = random.Random(123456)
+    high = 200
+    rng = random.Random(123)
     gen = stable_random_distribution(high, rng=rng)
     lst = islice(gen, 1000)
 
@@ -154,4 +154,4 @@ def test_fill_domain_speed():
         if stable_steps < uniform_steps:
             wins += 1
 
-    assert wins / trials > 0.90
+    assert wins / trials > 0.999

micall/utils/stable_random_distribution.py

@@ -1,10 +1,8 @@
-from typing import Iterator, Optional
+from typing import Iterator, Optional, Sequence
 
 import random
 import numpy as np
 
-DUPLICATION_FACTOR = 1
-
 
 def stable_random_distribution(high: int,
                                rng: Optional[random.Random] = None,
@@ -16,14 +14,13 @@ def stable_random_distribution(high: int,
     if rng is None:
         rng = random.Random()
 
-    maximum = high - 1
-    block = np.arange(high)
-    population = np.concatenate([block] * DUPLICATION_FACTOR, axis=0)
-
-    assert len(population) == DUPLICATION_FACTOR * len(block)
+    population = np.arange(high)
+    weights = np.zeros(high) + 1
 
     while True:
-        choice = rng.randint(0, maximum)
-        index = population[choice]
+        pweights: Sequence[float] = weights  # type: ignore
+        index = rng.choices(population, weights=pweights)[0]
         yield index
-        population[choice] = rng.randint(0, maximum)
+        weights[index] /= 2
+        if weights[index] < 0.01:
+            weights += 1