A weighted version of random.choice

后端 未结 25 1922
闹比i
闹比i 2020-11-21 06:29

I needed to write a weighted version of random.choice (each element in the list has a different probability for being selected). This is what I came up with:



        
25条回答
  •  时光说笑
    2020-11-21 06:47

    I didn't love the syntax of any of those. I really wanted to just specify what the items were and what the weighting of each was. I realize I could have used random.choices but instead I quickly wrote the class below.

    import random, string
    from numpy import cumsum
    
    class randomChoiceWithProportions:
        '''
        Accepts a dictionary of choices as keys and weights as values. Example if you want a unfair dice:
    
    
        choiceWeightDic = {"1":0.16666666666666666, "2": 0.16666666666666666, "3": 0.16666666666666666
        , "4": 0.16666666666666666, "5": .06666666666666666, "6": 0.26666666666666666}
        dice = randomChoiceWithProportions(choiceWeightDic)
    
        samples = []
        for i in range(100000):
            samples.append(dice.sample())
    
        # Should be close to .26666
        samples.count("6")/len(samples)
    
        # Should be close to .16666
        samples.count("1")/len(samples)
        '''
        def __init__(self, choiceWeightDic):
            self.choiceWeightDic = choiceWeightDic
            weightSum = sum(self.choiceWeightDic.values())
            assert weightSum == 1, 'Weights sum to ' + str(weightSum) + ', not 1.'
            self.valWeightDict = self._compute_valWeights()
    
        def _compute_valWeights(self):
            valWeights = list(cumsum(list(self.choiceWeightDic.values())))
            valWeightDict = dict(zip(list(self.choiceWeightDic.keys()), valWeights))
            return valWeightDict
    
        def sample(self):
            num = random.uniform(0,1)
            for key, val in self.valWeightDict.items():
                if val >= num:
                    return key
    

提交回复
热议问题