Moved random_weighted_selection to utils and rewrote it using bisect.

darius · darius · commit dafac5b12a1c · 2011-10-29T20:05:31.000Z
diff --git a/search.py b/search.py
@@ -367,7 +367,8 @@ def genetic_algorithm(population, fitness_fn, ngen=1000, pmut=0.1):
     for i in range(ngen):
         new_population = []
         for i in len(population):
-            p1, p2 = random_weighted_selections(population, 2, fitness_fn)
+            fitnesses = map(fitness_fn, population)
+            p1, p2 = weighted_sample_with_replacement(population, fitnesses, 2)
             child = p1.mate(p2)
             if random.uniform(0, 1) < pmut:
                 child.mutate()
@@ -389,25 +390,6 @@ def mutate(self):
         "Change a few of my genes."
         abstract
 
-def random_weighted_selection(seq, n, weight_fn):
-    """Pick n elements of seq, weighted according to weight_fn.
-    That is, apply weight_fn to each element of seq, add up the total.
-    Then choose an element e with probability weight[e]/total.
-    Repeat n times, with replacement. """
-    totals = []; runningtotal = 0
-    for item in seq:
-        runningtotal += weight_fn(item)
-        totals.append(runningtotal)
-    selections = []
-    for s in range(n):
-        r = random.uniform(0, totals[-1])
-        for i in range(len(seq)):
-            if totals[i] > r:
-                selections.append(seq[i])
-                break
-    return selections
-
-
 #_____________________________________________________________________________
 # The remainder of this file implements examples for the search algorithms.
 
@@ -885,7 +867,4 @@ def compare_graph_searchers():
 
 >>> boggle_hill_climbing(list('ABCDEFGHI'), verbose=False)
 (['E', 'P', 'R', 'D', 'O', 'A', 'G', 'S', 'T'], 123)
-
->>> random_weighted_selection(range(10), 3, lambda x: x * x)
-[8, 9, 6]
 """)
diff --git a/utils.py b/utils.py
@@ -501,6 +501,18 @@ def probability(p):
     "Return true with probability p."
     return p > random.uniform(0.0, 1.0)
 
+def weighted_sample_with_replacement(seq, weights, n):
+    """Pick n samples from seq at random, with replacement, with the
+    probability of each element in proportion to its corresponding
+    weight."""
+    totals = []
+    for w in weights:
+        totals.append(w + totals[-1] if totals else w)
+    def sample():
+        r = random.uniform(0, totals[-1])
+        return seq[bisect.bisect(totals, r)]
+    return [sample() for s in range(n)]
+
 def num_or_str(x):
     """The argument is a string; convert to a number if possible, or strip it.
     >>> num_or_str('42')
@@ -939,4 +951,15 @@ def fixup(test):
 >>> sl(a)
 ['a', 'b', 'c', 'd', 'r', 'w', 'y', 'z']
 
+>>> weighted_sample_with_replacement([], [], 0)
+[]
+>>> weighted_sample_with_replacement('a', [3], 2)
+['a', 'a']
+>>> weighted_sample_with_replacement('ab', [0, 3], 3)
+['b', 'b', 'b']
 """
+
+__doc__ += random_tests("""
+>>> weighted_sample_with_replacement(range(10), [x*x for x in range(10)], 3)
+[8, 9, 6]
+""")
