bpo-40541: Add optional *counts* parameter to random.sample()

Doc/library/random.rst

@@ -217,7 +217,7 @@ Functions for sequences
       The optional parameter *random*.
 
 
-.. function:: sample(population, k)
+.. function:: sample(population, k, *, weights=None)
 
    Return a *k* length list of unique elements chosen from the population sequence
    or set. Used for random sampling without replacement.
@@ -231,13 +231,20 @@ Functions for sequences
    Members of the population need not be :term:`hashable` or unique.  If the population
    contains repeats, then each occurrence is a possible selection in the sample.
 
+   If *weights* are given, they must be non-negative integer counts.
+   Each selection effectively reduces the count by one, lowering
+   the probablity for the next selection.
+
    To choose a sample from a range of integers, use a :func:`range` object as an
    argument.  This is especially fast and space efficient for sampling from a large
    population:  ``sample(range(10000000), k=60)``.
 
    If the sample size is larger than the population size, a :exc:`ValueError`
    is raised.
 
+   .. versionchanged 3.9
+      Added the *weights* parameter.
+
    .. deprecated:: 3.9
       In the future, the *population* must be a sequence.  Instances of
       :class:`set` are no longer supported.  The set must first be converted
@@ -420,12 +427,11 @@ Simulations::
    >>> choices(['red', 'black', 'green'], [18, 18, 2], k=6)
    ['red', 'green', 'black', 'black', 'red', 'black']
 
-   >>> # Deal 20 cards without replacement from a deck of 52 playing cards
-   >>> # and determine the proportion of cards with a ten-value
-   >>> # (a ten, jack, queen, or king).
-   >>> deck = collections.Counter(tens=16, low_cards=36)
-   >>> seen = sample(list(deck.elements()), k=20)
-   >>> seen.count('tens') / 20
+   >>> # Deal 20 cards without replacement from a deck
+   >>> # of 52 playing cards, and determine the proportion of cards
+   >>> # with a ten-value:  ten, jack, queen, or king.
+   >>> dealt = sample(['tens', 'low cards'], weights=[16, 36], k=20)
+   >>> dealt.count('tens') / 20
    0.15
 
    >>> # Estimate the probability of getting 5 or more heads from 7 spins

Lib/random.py

@@ -331,7 +331,7 @@ def shuffle(self, x, random=None):
                 j = _int(random() * (i+1))
                 x[i], x[j] = x[j], x[i]
 
-    def sample(self, population, k):
+    def sample(self, population, k, *, weights=None):
         """Chooses k unique random elements from a population sequence or set.
 
         Returns a new list containing elements from the population while
@@ -340,6 +340,10 @@ def sample(self, population, k):
         samples.  This allows raffle winners (the sample) to be partitioned
         into grand prize and second place winners (the subslices).
 
+        If weights are given, they must be non-negative integer counts.
+        Each selection effectively reduces the count by one, lowering
+        the probablity for the next selection.
+
         Members of the population need not be hashable or unique.  If the
         population contains repeats, then each occurrence is a possible
         selection in the sample.
@@ -379,6 +383,16 @@ def sample(self, population, k):
             population = tuple(population)
         if not isinstance(population, _Sequence):
             raise TypeError("Population must be a sequence.  For dicts or sets, use sorted(d).")
+        if weights is not None:
+            cum_weights = list(_accumulate(weights))
+            total = cum_weights.pop()
+            if not isinstance(total, int):
+                raise TypeError('Weights must be integers')
+            if total < 0:
+                raise ValueError('Total of weights must be greater than zero')
+            selections = sample(range(total), k=k)
+            bisect = _bisect
+            return [population[bisect(cum_weights, s)] for s in selections]
         randbelow = self._randbelow
         n = len(population)
         if not 0 <= k <= n:

Misc/NEWS.d/next/Library/2020-05-06-15-36-47.bpo-40541.LlYghL.rst

@@ -0,0 +1 @@
+Added an optional *weights* parameter to random.sample().