Change locking boundaries for shared dict. Add a unit test.

[python_utils.git] / list_utils.py

diff --git a/list_utils.py b/list_utils.py
index 05512b564c2303372ff81660840bfd29dcba942b..b65ab0d69fab36196a94b4ef3327fae63b88f48d 100644 (file)
--- a/list_utils.py
+++ b/list_utils.py
@@ -2,7 +2,7 @@
 
 from collections import Counter
 from itertools import islice
-from typing import Any, Iterator, List, Mapping, Sequence
+from typing import Any, Iterator, List, Mapping, Sequence, Tuple
 
 
 def shard(lst: List[Any], size: int) -> Iterator[Any]:
@@ -197,9 +197,91 @@ 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]):
+    """
+    Returns all permutations of a sequence; takes O(N^2) time.
+
+    >>> for x in permute('cat'):
+    ...     print(x)
+    cat
+    cta
+    act
+    atc
+    tca
+    tac
+
+    """
+    yield from _permute(seq, "")
+
+
+def _permute(seq: Sequence[Any], path):
+    if len(seq) == 0:
+        yield path
+
+    for i in range(len(seq)):
+        car = seq[i]
+        left = seq[0:i]
+        right = seq[i + 1 :]
+        cdr = left + right
+        yield from _permute(cdr, path + car)
+
+
+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
+    target whether it was found or not.
+
+    >>> a = [1, 4, 5, 6, 7, 9, 10, 11]
+    >>> binary_search(a, 4)
+    (True, 1)
+
+    >>> binary_search(a, 12)
+    (False, 8)
+
+    >>> binary_search(a, 3)
+    (False, 1)
+
+    >>> 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)
+
+
+def _binary_search(
+    lst: Sequence[Any], target: Any, low: int, high: int
+) -> Tuple[bool, int]:
+    if high >= low:
+        mid = (high + low) // 2
+        if lst[mid] == target:
+            return (True, mid)
+        elif lst[mid] > target:
+            return _binary_search(lst, target, low, mid - 1)
+        else:
+            return _binary_search(lst, target, mid + 1, high)
+    else:
+        return (False, low)
 
 
 if __name__ == '__main__':
     import doctest
+
     doctest.testmod()