+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=1): + """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 + + # Fast path for weighted_choice() + u = _bisect(cumulative_dist, self.random() * total_sum) + if indices is not None: + u = indices[u] + k = 1.0 / total_sum + cumulative_dist = [k * s for s in cumulative_dist] + results = [] + if indices is None: + for i in range(0,amount): + u = self.random() + results.append(_bisect(cumulative_dist, u)) + return results + else: + for i in range(0,amount): + u = self.random() + results.append(indices[_bisect(cumulative_dist, u)]) + return results +weighted_choice = _inst.weighted_choice