Adds shuffle/scramble to list_utils.

[python_utils.git] / list_utils.py

diff --git a/list_utils.py b/list_utils.py
index 88c436be24f88e73093f2c6f509a157e181ed6f9..7f28ade922a4b733ee29ea0c68bc957ff587c51b 100644 (file)
--- a/list_utils.py
+++ b/list_utils.py
@@ -1,8 +1,11 @@
 #!/usr/bin/env python3
 
+"""Some useful(?) utilities for dealing with Lists."""
+
+import random
 from collections import Counter
 from itertools import islice
-from typing import Any, Iterator, List, Mapping, Sequence, Tuple
+from typing import Any, Iterator, List, MutableSequence, Sequence, Tuple
 
 
 def shard(lst: List[Any], size: int) -> Iterator[Any]:
@@ -65,7 +68,7 @@ def remove_list_if_one_element(lst: List[Any]) -> Any:
         return lst
 
 
-def population_counts(lst: List[Any]) -> Mapping[Any, int]:
+def population_counts(lst: Sequence[Any]) -> Counter:
     """
     Return a population count mapping for the list (i.e. the keys are
     list items and the values are the number of occurrances of that
@@ -197,12 +200,12 @@ def ngrams(lst: Sequence[Any], n):
     ['an', 'awesome', 'test']
     """
     for i in range(len(lst) - n + 1):
-        yield lst[i:i + n]
+        yield lst[i : i + n]
 
 
-def permute(seq: Sequence[Any]):
+def permute(seq: str):
     """
-    Returns all permutations of a sequence; takes O(N^2) time.
+    Returns all permutations of a sequence; takes O(N!) time.
 
     >>> for x in permute('cat'):
     ...     print(x)
@@ -217,19 +220,44 @@ def permute(seq: Sequence[Any]):
     yield from _permute(seq, "")
 
 
-def _permute(seq: Sequence[Any], path):
-    if len(seq) == 0:
+def _permute(seq: str, path: str):
+    seq_len = len(seq)
+    if seq_len == 0:
         yield path
 
-    for i in range(len(seq)):
+    for i in range(seq_len):
         car = seq[i]
         left = seq[0:i]
-        right = seq[i + 1:]
+        right = seq[i + 1 :]
         cdr = left + right
         yield from _permute(cdr, path + car)
 
 
-def binary_search(lst: Sequence[Any], target: Any) -> Tuple[bool, int]:
+def shuffle(seq: MutableSequence[Any]) -> MutableSequence[Any]:
+    """Shuffles a sequence into a random order.
+
+    >>> random.seed(22)
+    >>> shuffle([1, 2, 3, 4, 5])
+    [3, 4, 1, 5, 2]
+
+    >>> shuffle('example')
+    'empaelx'
+
+    """
+    if isinstance(seq, str):
+        import string_utils
+
+        return string_utils.shuffle(seq)
+    else:
+        random.shuffle(seq)
+        return seq
+
+
+def scramble(seq: MutableSequence[Any]) -> MutableSequence[Any]:
+    return shuffle(seq)
+
+
+def binary_search(lst: Sequence[Any], target: Any, *, sanity_check=False) -> Tuple[bool, int]:
     """Performs a binary search on lst (which must already be sorted).
     Returns a Tuple composed of a bool which indicates whether the
     target was found and an int which indicates the index closest to
@@ -248,7 +276,19 @@ def binary_search(lst: Sequence[Any], target: Any) -> Tuple[bool, int]:
     >>> binary_search(a, 2)
     (False, 1)
 
+    >>> a.append(9)
+    >>> binary_search(a, 4, sanity_check=True)
+    Traceback (most recent call last):
+    ...
+    AssertionError
+
     """
+    if sanity_check:
+        last = None
+        for x in lst:
+            if last is not None:
+                assert x >= last  # This asserts iff the list isn't sorted
+            last = x  # in ascending order.
     return _binary_search(lst, target, 0, len(lst) - 1)
 
 
@@ -267,4 +307,5 @@ def _binary_search(lst: Sequence[Any], target: Any, low: int, high: int) -> Tupl
 
 if __name__ == '__main__':
     import doctest
+
     doctest.testmod()