A weighted version of random.choice

后端 未结 25 1921
闹比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:53

    Here's is the version that is being included in the standard library for Python 3.6:

    import itertools as _itertools
    import bisect as _bisect
    
    class Random36(random.Random):
        "Show the code included in the Python 3.6 version of the Random class"
    
        def choices(self, population, weights=None, *, cum_weights=None, k=1):
            """Return a k sized list of population elements chosen with replacement.
    
            If the relative weights or cumulative weights are not specified,
            the selections are made with equal probability.
    
            """
            random = self.random
            if cum_weights is None:
                if weights is None:
                    _int = int
                    total = len(population)
                    return [population[_int(random() * total)] for i in range(k)]
                cum_weights = list(_itertools.accumulate(weights))
            elif weights is not None:
                raise TypeError('Cannot specify both weights and cumulative weights')
            if len(cum_weights) != len(population):
                raise ValueError('The number of weights does not match the population')
            bisect = _bisect.bisect
            total = cum_weights[-1]
            return [population[bisect(cum_weights, random() * total)] for i in range(k)]
    

    Source: https://hg.python.org/cpython/file/tip/Lib/random.py#l340

提交回复
热议问题