+from _collections_abc import Set as _Set, Sequence as _Sequence, Mapping as _Mapping +from itertools import accumulate as _accumulate +from bisect import bisect as _bisect + "SystemRandom", "weighted_choice"] + def weighted_choice(self, data, amount=None): + """An iterator of random elements with the chances defined by relative + weights. + + If argument is a mapping then generates random keys with the + probability that is proportional to the value mapped from this key. + + If argument is a sequence then generates random indices with the + probability that is proportional to value at this index. + """ + if isinstance(data, _Mapping): + indices = list(data.keys()) + weights = data.values() + else: + indices = None + weights = data + cumulative_dist = list(_accumulate(weights)) + total_sum = cumulative_dist[-1] + if any(w < 0 for w in weights): + raise ValueError("All weights must be non-negative") + if not total_sum: + raise ValueError("At least one weight must be greater than zero") + del weights + k = 1.0 / total_sum + cumulative_dist = [k * s for s in cumulative_dist] + results = [] + if amount is None: + u = self.random() + if indices is None: + return _bisect(cumulative_dist, u) + else: + return indices[_bisect(cumulative_dist, u)] + else: + if amount < 0: + raise ValueError("Must request at least one value") + for i in range(0,amount): + u = self.random() + if indices is None: + results.append(_bisect(cumulative_dist, u)) + else: + results.append(indices[_bisect(cumulative_dist, u)]) + return results +weighted_choice = _inst.weighted_choice