recipes.py 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242
  1. from itertools import *
  2. from collections import *
  3. def take(n, iterable):
  4. "Return first n items of the iterable as a list"
  5. return list(islice(iterable, n))
  6. def prepend(value, iterator):
  7. "Prepend a single value in front of an iterator"
  8. # prepend(1, [2, 3, 4]) -> 1 2 3 4
  9. return chain([value], iterator)
  10. def tabulate(function, start=0):
  11. "Return function(0), function(1), ..."
  12. return map(function, count(start))
  13. def tail(n, iterable):
  14. "Return an iterator over the last n items"
  15. # tail(3, 'ABCDEFG') --> E F G
  16. return iter(collections.deque(iterable, maxlen=n))
  17. def consume(iterator, n=None):
  18. "Advance the iterator n-steps ahead. If n is None, consume entirely."
  19. # Use functions that consume iterators at C speed.
  20. if n is None:
  21. # feed the entire iterator into a zero-length deque
  22. collections.deque(iterator, maxlen=0)
  23. else:
  24. # advance to the empty slice starting at position n
  25. next(islice(iterator, n, n), None)
  26. def nth(iterable, n, default=None):
  27. "Returns the nth item or a default value"
  28. return next(islice(iterable, n, None), default)
  29. def all_equal(iterable):
  30. "Returns True if all the elements are equal to each other"
  31. g = groupby(iterable)
  32. return next(g, True) and not next(g, False)
  33. def quantify(iterable, pred=bool):
  34. "Count how many times the predicate is true"
  35. return sum(map(pred, iterable))
  36. def padnone(iterable):
  37. """Returns the sequence elements and then returns None indefinitely.
  38. Useful for emulating the behavior of the built-in map() function.
  39. """
  40. return chain(iterable, repeat(None))
  41. def ncycles(iterable, n):
  42. "Returns the sequence elements n times"
  43. return chain.from_iterable(repeat(tuple(iterable), n))
  44. def dotproduct(vec1, vec2):
  45. return sum(map(operator.mul, vec1, vec2))
  46. def flatten(listOfLists):
  47. "Flatten one level of nesting"
  48. return chain.from_iterable(listOfLists)
  49. def repeatfunc(func, times=None, *args):
  50. """Repeat calls to func with specified arguments.
  51. Example: repeatfunc(random.random)
  52. """
  53. if times is None:
  54. return starmap(func, repeat(args))
  55. return starmap(func, repeat(args, times))
  56. def pairwise(iterable):
  57. "s -> (s0,s1), (s1,s2), (s2, s3), ..."
  58. a, b = tee(iterable)
  59. next(b, None)
  60. return zip(a, b)
  61. def grouper(iterable, n, fillvalue=None):
  62. "Collect data into fixed-length chunks or blocks"
  63. # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
  64. args = [iter(iterable)] * n
  65. return zip_longest(*args, fillvalue=fillvalue)
  66. def roundrobin(*iterables):
  67. "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
  68. # Recipe credited to George Sakkis
  69. num_active = len(iterables)
  70. nexts = cycle(iter(it).__next__ for it in iterables)
  71. while num_active:
  72. try:
  73. for next in nexts:
  74. yield next()
  75. except StopIteration:
  76. # Remove the iterator we just exhausted from the cycle.
  77. num_active -= 1
  78. nexts = cycle(islice(nexts, num_active))
  79. def partition(pred, iterable):
  80. 'Use a predicate to partition entries into false entries and true entries'
  81. # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9
  82. t1, t2 = tee(iterable)
  83. return filterfalse(pred, t1), filter(pred, t2)
  84. def powerset(iterable):
  85. "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
  86. s = list(iterable)
  87. return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
  88. def unique_everseen(iterable, key=None):
  89. "List unique elements, preserving order. Remember all elements ever seen."
  90. # unique_everseen('AAAABBBCCDAABBB') --> A B C D
  91. # unique_everseen('ABBCcAD', str.lower) --> A B C D
  92. seen = set()
  93. seen_add = seen.add
  94. if key is None:
  95. for element in filterfalse(seen.__contains__, iterable):
  96. seen_add(element)
  97. yield element
  98. else:
  99. for element in iterable:
  100. k = key(element)
  101. if k not in seen:
  102. seen_add(k)
  103. yield element
  104. def unique_justseen(iterable, key=None):
  105. "List unique elements, preserving order. Remember only the element just seen."
  106. # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
  107. # unique_justseen('ABBCcAD', str.lower) --> A B C A D
  108. return map(next, map(itemgetter(1), groupby(iterable, key)))
  109. def iter_except(func, exception, first=None):
  110. """ Call a function repeatedly until an exception is raised.
  111. Converts a call-until-exception interface to an iterator interface.
  112. Like builtins.iter(func, sentinel) but uses an exception instead
  113. of a sentinel to end the loop.
  114. Examples:
  115. iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator
  116. iter_except(d.popitem, KeyError) # non-blocking dict iterator
  117. iter_except(d.popleft, IndexError) # non-blocking deque iterator
  118. iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue
  119. iter_except(s.pop, KeyError) # non-blocking set iterator
  120. """
  121. try:
  122. if first is not None:
  123. yield first() # For database APIs needing an initial cast to db.first()
  124. while True:
  125. yield func()
  126. except exception:
  127. pass
  128. def first_true(iterable, default=False, pred=None):
  129. """Returns the first true value in the iterable.
  130. If no true value is found, returns *default*
  131. If *pred* is not None, returns the first item
  132. for which pred(item) is true.
  133. """
  134. # first_true([a,b,c], x) --> a or b or c or x
  135. # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
  136. return next(filter(pred, iterable), default)
  137. def random_product(*args, repeat=1):
  138. "Random selection from itertools.product(*args, **kwds)"
  139. pools = [tuple(pool) for pool in args] * repeat
  140. return tuple(random.choice(pool) for pool in pools)
  141. def random_permutation(iterable, r=None):
  142. "Random selection from itertools.permutations(iterable, r)"
  143. pool = tuple(iterable)
  144. r = len(pool) if r is None else r
  145. return tuple(random.sample(pool, r))
  146. def random_combination(iterable, r):
  147. "Random selection from itertools.combinations(iterable, r)"
  148. pool = tuple(iterable)
  149. n = len(pool)
  150. indices = sorted(random.sample(range(n), r))
  151. return tuple(pool[i] for i in indices)
  152. def random_combination_with_replacement(iterable, r):
  153. "Random selection from itertools.combinations_with_replacement(iterable, r)"
  154. pool = tuple(iterable)
  155. n = len(pool)
  156. indices = sorted(random.randrange(n) for i in range(r))
  157. return tuple(pool[i] for i in indices)
  158. def nth_combination(iterable, r, index):
  159. 'Equivalent to list(combinations(iterable, r))[index]'
  160. pool = tuple(iterable)
  161. n = len(pool)
  162. if r < 0 or r > n:
  163. raise ValueError
  164. c = 1
  165. k = min(r, n-r)
  166. for i in range(1, k+1):
  167. c = c * (n - k + i) // i
  168. if index < 0:
  169. index += c
  170. if index < 0 or index >= c:
  171. raise IndexError
  172. result = []
  173. while r:
  174. c, n, r = c*r//n, n-1, r-1
  175. while index >= c:
  176. index -= c
  177. c, n = c*(n-r)//n, n-1
  178. result.append(pool[-1-n])
  179. return tuple(result)