3 """Some useful(?) utilities for dealing with Lists."""
6 from collections import Counter
7 from itertools import islice
8 from typing import Any, Iterator, List, MutableSequence, Sequence, Tuple
11 def shard(lst: List[Any], size: int) -> Iterator[Any]:
13 Yield successive size-sized shards from lst.
15 >>> for sublist in shard([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], 3):
16 ... [_ for _ in sublist]
23 for x in range(0, len(lst), size):
24 yield islice(lst, x, x + size)
27 def flatten(lst: List[Any]) -> List[Any]:
31 >>> flatten([ 1, [2, 3, 4, [5], 6], 7, [8, [9]]])
32 [1, 2, 3, 4, 5, 6, 7, 8, 9]
37 if isinstance(lst[0], list):
38 return flatten(lst[0]) + flatten(lst[1:])
39 return lst[:1] + flatten(lst[1:])
42 def prepend(item: Any, lst: List[Any]) -> List[Any]:
44 Prepend an item to a list.
46 >>> prepend('foo', ['bar', 'baz'])
54 def remove_list_if_one_element(lst: List[Any]) -> Any:
56 Remove the list and return the 0th element iff its length is one.
58 >>> remove_list_if_one_element([1234])
61 >>> remove_list_if_one_element([1, 2, 3, 4])
71 def population_counts(lst: Sequence[Any]) -> Counter:
73 Return a population count mapping for the list (i.e. the keys are
74 list items and the values are the number of occurrances of that
75 list item in the original list.
77 >>> population_counts([1, 1, 1, 2, 2, 3, 3, 3, 4])
78 Counter({1: 3, 3: 3, 2: 2, 4: 1})
84 def most_common(lst: List[Any], *, count=1) -> Any:
87 Return the most common item in the list. In the case of ties,
88 which most common item is returned will be random.
90 >>> most_common([1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4])
93 >>> most_common([1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4], count=2)
97 p = population_counts(lst)
98 return remove_list_if_one_element([_[0] for _ in p.most_common()[0:count]])
101 def least_common(lst: List[Any], *, count=1) -> Any:
103 Return the least common item in the list. In the case of
104 ties, which least common item is returned will be random.
106 >>> least_common([1, 1, 1, 2, 2, 3, 3, 3, 4])
109 >>> least_common([1, 1, 1, 2, 2, 3, 3, 3, 4], count=2)
113 p = population_counts(lst)
114 mc = p.most_common()[-count:]
116 return remove_list_if_one_element([_[0] for _ in mc])
119 def dedup_list(lst: List[Any]) -> List[Any]:
121 Remove duplicates from the list performantly.
123 >>> dedup_list([1, 2, 1, 3, 3, 4, 2, 3, 4, 5, 1])
127 return list(set(lst))
130 def uniq(lst: List[Any]) -> List[Any]:
132 Alias for dedup_list.
134 return dedup_list(lst)
137 def contains_duplicates(lst: List[Any]) -> bool:
139 Does the list contian duplicate elements or not?
141 >>> lst = [1, 2, 1, 3, 3, 4, 4, 5, 6, 1, 3, 4]
142 >>> contains_duplicates(lst)
145 >>> contains_duplicates(dedup_list(lst))
157 def all_unique(lst: List[Any]) -> bool:
159 Inverted alias for contains_duplicates.
161 return not contains_duplicates(lst)
164 def transpose(lst: List[Any]) -> List[Any]:
166 Transpose a list of lists.
168 >>> lst = [[1, 2], [3, 4], [5, 6]]
170 [[1, 3, 5], [2, 4, 6]]
173 transposed = zip(*lst)
174 return [list(_) for _ in transposed]
177 def ngrams(lst: Sequence[Any], n):
179 Return the ngrams in the sequence.
181 >>> seq = 'encyclopedia'
182 >>> for _ in ngrams(seq, 3):
195 >>> seq = ['this', 'is', 'an', 'awesome', 'test']
196 >>> for _ in ngrams(seq, 3):
199 ['is', 'an', 'awesome']
200 ['an', 'awesome', 'test']
202 for i in range(len(lst) - n + 1):
206 def permute(seq: str):
208 Returns all permutations of a sequence; takes O(N!) time.
210 >>> for x in permute('cat'):
220 yield from _permute(seq, "")
223 def _permute(seq: str, path: str):
228 for i in range(seq_len):
233 yield from _permute(cdr, path + car)
236 def shuffle(seq: MutableSequence[Any]) -> MutableSequence[Any]:
237 """Shuffles a sequence into a random order.
240 >>> shuffle([1, 2, 3, 4, 5])
243 >>> shuffle('example')
247 if isinstance(seq, str):
250 return string_utils.shuffle(seq)
256 def scramble(seq: MutableSequence[Any]) -> MutableSequence[Any]:
260 def binary_search(lst: Sequence[Any], target: Any, *, sanity_check=False) -> Tuple[bool, int]:
261 """Performs a binary search on lst (which must already be sorted).
262 Returns a Tuple composed of a bool which indicates whether the
263 target was found and an int which indicates the index closest to
264 target whether it was found or not.
266 >>> a = [1, 4, 5, 6, 7, 9, 10, 11]
267 >>> binary_search(a, 4)
270 >>> binary_search(a, 12)
273 >>> binary_search(a, 3)
276 >>> binary_search(a, 2)
280 >>> binary_search(a, 4, sanity_check=True)
281 Traceback (most recent call last):
290 assert x >= last # This asserts iff the list isn't sorted
291 last = x # in ascending order.
292 return _binary_search(lst, target, 0, len(lst) - 1)
295 def _binary_search(lst: Sequence[Any], target: Any, low: int, high: int) -> Tuple[bool, int]:
297 mid = (high + low) // 2
298 if lst[mid] == target:
300 elif lst[mid] > target:
301 return _binary_search(lst, target, low, mid - 1)
303 return _binary_search(lst, target, mid + 1, high)
308 if __name__ == '__main__':