X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=list_utils.py;h=91af8f9eb924fb7e7e04932d58a1bcb6eded0690;hb=e8fbbb7306430478dec55d2c963eed116d8330cc;hp=182e2bc5c104908f39a15e4675021e6ed8a7c338;hpb=986d5f7ada15e56019518db43d07b76f94468e1a;p=python_utils.git diff --git a/list_utils.py b/list_utils.py index 182e2bc..91af8f9 100644 --- a/list_utils.py +++ b/list_utils.py @@ -1,8 +1,10 @@ #!/usr/bin/env python3 +"""Some useful(?) utilities for dealing with Lists.""" + from collections import Counter from itertools import islice -from typing import Any, Iterator, List, Mapping, Sequence +from typing import Any, Iterator, List, Sequence, Tuple def shard(lst: List[Any], size: int) -> Iterator[Any]: @@ -48,7 +50,24 @@ def prepend(item: Any, lst: List[Any]) -> List[Any]: return lst -def population_counts(lst: List[Any]) -> Mapping[Any, int]: +def remove_list_if_one_element(lst: List[Any]) -> Any: + """ + Remove the list and return the 0th element iff its length is one. + + >>> remove_list_if_one_element([1234]) + 1234 + + >>> remove_list_if_one_element([1, 2, 3, 4]) + [1, 2, 3, 4] + + """ + if len(lst) == 1: + return lst[0] + else: + return lst + + +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 @@ -61,29 +80,39 @@ def population_counts(lst: List[Any]) -> Mapping[Any, int]: return Counter(lst) -def most_common_item(lst: List[Any]) -> Any: +def most_common(lst: List[Any], *, count=1) -> Any: """ Return the most common item in the list. In the case of ties, which most common item is returned will be random. - >>> most_common_item([1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4]) + >>> most_common([1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4]) 3 + >>> most_common([1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4], count=2) + [3, 1] + """ - return population_counts(lst).most_common(1)[0][0] + p = population_counts(lst) + return remove_list_if_one_element([_[0] for _ in p.most_common()[0:count]]) -def least_common_item(lst: List[Any]) -> Any: +def least_common(lst: List[Any], *, count=1) -> Any: """ Return the least common item in the list. In the case of ties, which least common item is returned will be random. - >>> least_common_item([1, 1, 1, 2, 2, 3, 3, 3, 4]) + >>> least_common([1, 1, 1, 2, 2, 3, 3, 3, 4]) 4 + >>> least_common([1, 1, 1, 2, 2, 3, 3, 3, 4], count=2) + [4, 2] + """ - return population_counts(lst).most_common()[-1][0] + p = population_counts(lst) + mc = p.most_common()[-count:] + mc.reverse() + return remove_list_if_one_element([_[0] for _ in mc]) def dedup_list(lst: List[Any]) -> List[Any]: @@ -100,16 +129,158 @@ def dedup_list(lst: List[Any]) -> List[Any]: def uniq(lst: List[Any]) -> List[Any]: """ Alias for dedup_list. - """ return dedup_list(lst) +def contains_duplicates(lst: List[Any]) -> bool: + """ + Does the list contian duplicate elements or not? + + >>> lst = [1, 2, 1, 3, 3, 4, 4, 5, 6, 1, 3, 4] + >>> contains_duplicates(lst) + True + + >>> contains_duplicates(dedup_list(lst)) + False + + """ + seen = set() + for _ in lst: + if _ in seen: + return True + seen.add(_) + return False + + +def all_unique(lst: List[Any]) -> bool: + """ + Inverted alias for contains_duplicates. + """ + return not contains_duplicates(lst) + + +def transpose(lst: List[Any]) -> List[Any]: + """ + Transpose a list of lists. + + >>> lst = [[1, 2], [3, 4], [5, 6]] + >>> transpose(lst) + [[1, 3, 5], [2, 4, 6]] + + """ + transposed = zip(*lst) + return [list(_) for _ in transposed] + + def ngrams(lst: Sequence[Any], n): + """ + Return the ngrams in the sequence. + + >>> seq = 'encyclopedia' + >>> for _ in ngrams(seq, 3): + ... _ + 'enc' + 'ncy' + 'cyc' + 'ycl' + 'clo' + 'lop' + 'ope' + 'ped' + 'edi' + 'dia' + + >>> seq = ['this', 'is', 'an', 'awesome', 'test'] + >>> for _ in ngrams(seq, 3): + ... _ + ['this', 'is', 'an'] + ['is', 'an', 'awesome'] + ['an', 'awesome', 'test'] + """ for i in range(len(lst) - n + 1): - yield lst[i:i + n] + yield lst[i : i + n] + + +def permute(seq: str): + """ + Returns all permutations of a sequence; takes O(N!) time. + + >>> for x in permute('cat'): + ... print(x) + cat + cta + act + atc + tca + tac + + """ + yield from _permute(seq, "") + + +def _permute(seq: str, path: str): + seq_len = len(seq) + if seq_len == 0: + yield path + + for i in range(seq_len): + 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()