Browse Source

added arg for pre-computed sum for weightedChoice

Darren Ranalli 21 years ago
parent
commit
f919dc99ab

+ 5 - 0
direct/src/directutil/WeightedChoice.py

@@ -1,6 +1,11 @@
 
 import random
 
+# DCR: I added a weightedChoice() function to PythonUtil that supports
+# floating-point weights and is intended for one-shot choices. It
+# has an optional 'sum' argument that you can pass in if you know the
+# sum of the weights and want to make repeated choices.
+
 class WeightedChoice:
     def __init__(self, listOfLists, weightIndex=0):
         t=0

+ 2 - 4
direct/src/level/DistributedLevelAI.py

@@ -6,8 +6,7 @@ import DistributedObjectAI
 import Level
 import DirectNotifyGlobal
 import EntityCreatorAI
-import WeightedChoice
-from PythonUtil import Functor
+from PythonUtil import Functor, weightedChoice
 
 class DistributedLevelAI(DistributedObjectAI.DistributedObjectAI,
                          Level.Level):
@@ -74,8 +73,7 @@ class DistributedLevelAI(DistributedObjectAI.DistributedObjectAI,
         # make list of lists: [(weight, scenarioIndex), ...]
         lol = zip([1] * levelSpec.getNumScenarios(),
                   range(levelSpec.getNumScenarios()))
-        wc = WeightedChoice.WeightedChoice(lol)
-        scenarioIndex = wc.choose()[1]
+        scenarioIndex = weightedChoice(lol)
 
         Level.Level.initializeLevel(self, self.doId, levelSpec, scenarioIndex)
 

+ 14 - 10
direct/src/showbase/PythonUtil.py

@@ -841,20 +841,24 @@ def mostDerivedLast(classList):
         return result
     classList.sort(compare)
 
-def weightedChoice(choiceList, rng=random.random):
-    """given a list of (probability,item) pairs, chooses an item based on the
-    probabilities. rng must return 0..1"""
-    sum = 0.
-    for prob, item in choiceList:
-        sum += prob
+def weightedChoice(choiceList, rng=random.random, sum=None):
+    """given a list of (weight,item) pairs, chooses an item based on the
+    weights. rng must return 0..1. if you happen to have the sum of the
+    weights, pass it in 'sum'."""
+    if sum is None:
+        sum = 0.
+        for weight, item in choiceList:
+            sum += weight
+
     rand = rng()
     accum = rand * sum
-    for prob, item in choiceList:
-        accum -= prob
+    for weight, item in choiceList:
+        accum -= weight
         if accum <= 0.:
             return item
-    # rand must be ~1., and floating-point error prevented accum from
-    # hitting 0. Return the last item.
+    # rand is ~1., and floating-point error prevented accum from hitting 0.
+    # Or you passed in a 'sum' that was was too large.
+    # Return the last item.
     return item
 
 def randFloat(a, b, rng=random.random):