Repository: more-itertools/more-itertools (python)

Issue:
Issues with sample()
I've only just started looking at this and have a few thoughts.

The comparison to the standard library should mention:

* The `counts` parameter in the `random.sample()` is not comparable to `weights` in `more_itertools.sample()`
The former allows an element to be selected multiple times to model classic "urn problems".  The latter just affects the likelihood of an element being selected at all.

* The standard library `random.sample()` raises an exception if *k* is larger than the population.  However `more_itertools.sample()` just returns the population unchanged and in order:

```
>>> import more_itertools
>>> more_itertools.sample(range(10), k=15)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

>>> import random
>>> random.sample(range(10), k=15)
Traceback (most recent call last):
   ...
ValueError: Sample larger than population or is negative
```
* The standard library `random.sample()` guarantees the results to be fully shuffled.  In contrast, the `more_itertools.sample()` initial reservoir fill is not shuffled at all:

```
>>> import more_itertools
>>> import random

>>> random.seed(8675309)
>>> more_itertools.sample(range(2000), k=1990)[:15]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]

>>> random.seed(8675309)
>>> random.sample(range(2000), k=1990)[:15]
[823, 438, 1044, 1923, 1359, 1032, 1762, 1838, 575, 1764, 1313, 1674, 1325, 1386, 465]
```

It would be great if a `random.shuffle()` step were added at the end. If not, then a `sorted()` step would be an improvement over having a result that is neither randomized or sorted.

Pull request description:
Issue 854: sample improvements
This PR makes some changes to sample:
* `shuffle` is called before returning the output
* A note about differences from `random.sample` is added to the docstring
* `ValueError` is raised for negative `k` values

Closes #854

Patch:
diff --git a/more_itertools/more.py b/more_itertools/more.py
index 7b48190..7922eb0 100755
--- a/more_itertools/more.py
+++ b/more_itertools/more.py
@@ -4,7 +4,7 @@ import warnings
 from collections import Counter, defaultdict, deque, abc
 from collections.abc import Sequence
 from functools import cached_property, partial, reduce, wraps
-from heapq import heapify, heapreplace, heappop
+from heapq import heapify, heapreplace
 from itertools import (
     chain,
     combinations,
@@ -23,7 +23,7 @@ from itertools import (
 )
 from math import comb, e, exp, factorial, floor, fsum, log, perm, tau
 from queue import Empty, Queue
-from random import random, randrange, uniform
+from random import random, randrange, shuffle, uniform
 from operator import itemgetter, mul, sub, gt, lt, ge, le
 from sys import hexversion, maxsize
 from time import monotonic
@@ -3535,12 +3535,14 @@ def map_if(iterable, pred, func, func_else=lambda x: x):
         yield func(item) if pred(item) else func_else(item)
 
 
-def _sample_unweighted(iterable, k):
+def _sample_unweighted(iterator, k, strict):
     # Implementation of "Algorithm L" from the 1994 paper by Kim-Hung Li:
     # "Reservoir-Sampling Algorithms of Time Complexity O(n(1+log(N/n)))".
 
     # Fill up the reservoir (collection of samples) with the first `k` samples
-    reservoir = take(k, iterable)
+    reservoir = take(k, iterator)
+    if strict and len(reservoir) < k:
+        raise ValueError('Sample larger than population')
 
     # Generate random number that's the largest in a sample of k U(0,1) numbers
     # Largest order statistic: https://en.wikipedia.org/wiki/Order_statistic
@@ -3550,17 +3552,18 @@ def _sample_unweighted(iterable, k):
     # number with a geometric distribution. Sample it using random() and logs.
     next_index = k + floor(log(random()) / log(1 - W))
 
-    for index, element in enumerate(iterable, k):
+    for index, element in enumerate(iterator, k):
         if index == next_index:
             reservoir[randrange(k)] = element
             # The new W is the largest in a sample of k U(0, `old_W`) numbers
             W *= exp(log(random()) / k)
             next_index += floor(log(random()) / log(1 - W)) + 1
 
+    shuffle(reservoir)
     return reservoir
 
 
-def _sample_weighted(iterable, k, weights):
+def _sample_weighted(iterator, k, weights, strict):
     # Implementation of "A-ExpJ" from the 2006 paper by Efraimidis et al. :
     # "Weighted random sampling with a reservoir".
 
@@ -3569,7 +3572,10 @@ def _sample_weighted(iterable, k, weights):
 
     # Fill up the reservoir (collection of samples) with the first `k`
     # weight-keys and elements, then heapify the list.
-    reservoir = take(k, zip(weight_keys, iterable))
+    reservoir = take(k, zip(weight_keys, iterator))
+    if strict and len(reservoir) < k:
+        raise ValueError('Sample larger than population')
+
     heapify(reservoir)
 
     # The number of jumps before changing the reservoir is a random variable
@@ -3577,7 +3583,7 @@ def _sample_weighted(iterable, k, weights):
     smallest_weight_key, _ = reservoir[0]
     weights_to_skip = log(random()) / smallest_weight_key
 
-    for weight, element in zip(weights, iterable):
+    for weight, element in zip(weights, iterator):
         if weight >= weights_to_skip:
             # The notation here is consistent with the paper, but we store
             # the weight-keys in log-space for better numerical stability.
@@ -3591,14 +3597,15 @@ def _sample_weighted(iterable, k, weights):
         else:
             weights_to_skip -= weight
 
-    # Equivalent to [element for weight_key, element in sorted(reservoir)]
-    return [heappop(reservoir)[1] for _ in range(k)]
+    ret = [element for weight_key, element in reservoir]
+    shuffle(ret)
+    return ret
 
 
-def sample(iterable, k, weights=None):
+def sample(iterable, k, weights=None, strict=False):
     """Return a *k*-length list of elements chosen (without replacement)
-    from the *iterable*. Like :func:`random.sample`, but works on iterables
-    of unknown length.
+    from the *iterable*. Similar to :func:`random.sample`, but works on
+    iterables of unknown length.
 
     >>> iterable = range(100)
     >>> sample(iterable, 5)  # doctest: +SKIP
@@ -3611,24 +3618,31 @@ def sample(iterable, k, weights=None):
     >>> sampled = sample(iterable, 5, weights=weights)  # doctest: +SKIP
     [79, 67, 74, 66, 78]
 
-    The algorithm can also be used to generate weighted random permutations.
-    The relative weight of each item determines the probability that it
-    appears late in the permutation.
+    Weighted selections are made without replacement.
+    After an element is selected, it is removed from the pool and the
+    relative weights of the other elements increase (this
+    does not match the behavior of :func:`random.sample`'s *counts*
+    parameter).
 
-    >>> data = "abcdefgh"
-    >>> weights = range(1, len(data) + 1)
-    >>> sample(data, k=len(data), weights=weights)  # doctest: +SKIP
-    ['c', 'a', 'b', 'e', 'g', 'd', 'h', 'f']
+    If the length of *iterable* is less than *k*,
+    ``ValueError`` is raised if *strict* is ``True`` and
+    all elements are returned (in shuffled order) if *strict* is ``False``.
+
+    By default, the `Algorithm L <https://w.wiki/ANrM>`__ reservoir sampling
+    technique is used. When *weights* are provided,
+    `Algorithm A-ExpJ <https://w.wiki/ANrS>`__ is used.
     """
+    if k < 0:
+        raise ValueError('k must be non-negative')
     if k == 0:
         return []
 
-    iterable = iter(iterable)
+    iterator = iter(iterable)
     if weights is None:
-        return _sample_unweighted(iterable, k)
+        return _sample_unweighted(iterator, k, strict)
     else:
         weights = iter(weights)
-        return _sample_weighted(iterable, k, weights)
+        return _sample_weighted(iterator, k, weights, strict)
 
 
 def is_sorted(iterable, key=None, reverse=False, strict=False):
diff --git a/more_itertools/more.pyi b/more_itertools/more.pyi
index e946023..7ddca59 100644
--- a/more_itertools/more.pyi
+++ b/more_itertools/more.pyi
@@ -538,10 +538,15 @@ def map_if(
     func: Callable[[Any], Any],
     func_else: Callable[[Any], Any] | None = ...,
 ) -> Iterator[Any]: ...
+def _sample_unweighted(iterator: Iterator[_T], k: int, strict) -> list[_T]: ...
+def _sample_weighted(
+    iterator: Iterator[_T], k: int, weights, strict
+) -> list[_T]: ...
 def sample(
     iterable: Iterable[_T],
     k: int,
     weights: Iterable[float] | None = ...,
+    strict: bool = False,
 ) -> list[_T]: ...
 def is_sorted(
     iterable: Iterable[_T],


Test patch:
diff --git a/tests/test_more.py b/tests/test_more.py
index 15efdaf..f451fc6 100644
--- a/tests/test_more.py
+++ b/tests/test_more.py
@@ -4145,6 +4145,11 @@ class SampleTests(TestCase):
         expected = ['f', 'e']
         self.assertEqual(actual, expected)
 
+    def test_negative(self):
+        data = [1, 2, 3, 4, 5]
+        with self.assertRaises(ValueError):
+            mi.sample(data, k=-1)
+
     def test_length(self):
         """Check that *k* elements are sampled."""
         data = [1, 2, 3, 4, 5]
@@ -4154,6 +4159,12 @@ class SampleTests(TestCase):
             expected = min(k, len(data))
             self.assertEqual(actual, expected)
 
+    def test_strict(self):
+        data = ['1', '2', '3', '4', '5']
+        self.assertEqual(set(mi.sample(data, 6, strict=False)), set(data))
+        with self.assertRaises(ValueError):
+            mi.sample(data, 6, strict=True)
+
     def test_sampling_entire_iterable(self):
         """If k=len(iterable), the sample contains the original elements."""
         data = ["a", 2, "a", 4, (1, 2, 3)]