#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
additions to itertools standard library
"""
__author__ = "Philippe Guglielmetti"
__copyright__ = "Copyright 2012, Philippe Guglielmetti"
__credits__ = ["functional toolset from http://pyeuler.wikidot.com/toolset",
"algos from https://github.com/tokland/pyeuler/blob/master/pyeuler/toolset.py",
"tools from http://docs.python.org/dev/py3k/library/itertools.html",
]
__license__ = "LGPL"
#!/usr/bin/python
from itertools import ifilter, islice, repeat, groupby
from itertools import count, imap, takewhile, tee, izip
from itertools import chain, starmap, cycle, dropwhile
import random
import logging
[docs]def take(n, iterable):
"""Take first n elements from iterable"""
return islice(iterable, n)
[docs]def index(n, iterable):
"Returns the nth item"
return islice(iterable, n, n+1).next()
[docs]def first(iterable):
"""Take first element in the iterable"""
return iterable.next()
[docs]def last(iterable):
"""Take last element in the iterable"""
return reduce(lambda x, y: y, iterable)
[docs]def take_every(n, iterable):
"""Take an element from iterator every n elements"""
return islice(iterable, 0, None, n)
[docs]def drop(n, iterable):
"""Drop n elements from iterable and return the rest"""
return islice(iterable, n, None)
[docs]def ilen(it):
"""Return length exhausing an iterator"""
return sum(1 for _ in it)
[docs]def irange(start_or_end, optional_end=None):
"""Return iterable that counts from start to end (both included)."""
if optional_end is None:
start, end = 0, start_or_end
else:
start, end = start_or_end, optional_end
return take(max(end - start + 1, 0), count(start))
[docs]def arange(start,stop,step=1.):
"""range for floats or other types"""
r = start
step=abs(step)
if stop<start :
while r > stop:
yield r
r -= step
else:
while r < stop:
yield r
r += step
[docs]def ilinear(start,end,n):
"""return iterator over n values linearly interpolated between (and including) start and end"""
if isinstance(start,(int,float)):
if start==end: #generate n times the same value for consistency
return repeat(start,n)
else: #make sure we generate n values including start and end
step=float(end-start)/(n-1)
return arange(start,end+step/2,step)
else: #suppose start and end are tuples or lists of the same size
res=(ilinear(s,e,n) for s,e in zip(start,end))
return izip(*res)
[docs]def flatten(lstlsts):
"""Flatten a list of lists"""
return (b for a in lstlsts for b in a)
[docs]def compact(it):
"""Filter None values from iterator"""
return ifilter(bool, it)
[docs]def groups(iterable, n, step):
"""Make groups of 'n' elements from the iterable advancing
'step' elements on each iteration"""
itlist = tee(iterable, n)
onestepit = izip(*(starmap(drop, enumerate(itlist))))
return take_every(step, onestepit)
[docs]def compose(f, g):
"""Compose two functions -> compose(f, g)(x) -> f(g(x))"""
def _wrapper(*args, **kwargs):
return f(g(*args, **kwargs))
return _wrapper
[docs]def iterate(func, arg):
"""After Haskell's iterate: apply function repeatedly."""
# not functional
while 1:
yield arg
arg = func(arg)
[docs]def tails(seq):
"""Get tails of a sequence: tails([1,2,3]) -> [1,2,3], [2,3], [3], []."""
for idx in xrange(len(seq)+1):
yield seq[idx:]
[docs]def ireduce(func, iterable, init=None):
"""Like reduce() but using iterators (a.k.a scanl)"""
# not functional
if init is None:
iterable = iter(iterable)
curr = iterable.next()
else:
curr = init
yield init
for x in iterable:
curr = func(curr, x)
yield curr
[docs]def unique(iterable, key=None):
"List unique elements, preserving order. Remember all elements ever seen."
# unique('AAAABBBCCDAABBB') --> A B C D
# unique('ABBCcAD', str.lower) --> A B C D
seen = set()
seen_add = seen.add
if key is None:
for element in iterable:
if element not in seen:
seen_add(element)
yield element
else:
for element in iterable:
k = key(element)
if k not in seen:
seen_add(k)
yield element
[docs]def identity(x):
"""Do nothing and return the variable untouched"""
return x
[docs]def occurrences(it, exchange=False):
"""Return dictionary with occurrences from iterable"""
return reduce(lambda occur, x: dict(occur, **{x: occur.get(x, 0) + 1}), it, {})
[docs]def product(*iterables, **kwargs):
"""http://stackoverflow.com/questions/12093364/cartesian-product-of-large-iterators-itertools"""
if len(iterables) == 0:
yield ()
else:
iterables = iterables * kwargs.get('repeat', 1)
it = iterables[0]
for item in it() if callable(it) else iter(it):
for items in product(*iterables[1:]):
yield (item, ) + items
# my functions added
[docs]def any(seq, pred=bool):
"Return True if pred(x) is True for at least one element in the iterable"
return (True in imap(pred, seq))
[docs]def all(seq, pred=bool):
"Return True if pred(x) is True for all elements in the iterable"
return (False not in imap(pred, seq))
[docs]def no(seq, pred=bool):
"Returns True if pred(x) is False for every element in the iterable"
return (True not in imap(pred, seq))
[docs]def takenth(n, iterable):
"Returns the nth item"
return islice(iterable, n, n+1).next()
[docs]def takeevery(n, iterable):
"""Take an element from iterator every n elements"""
return islice(iterable, 0, None, n)
[docs]def icross(*sequences):
"""Cartesian product of sequences (recursive version)"""
if sequences:
for x in sequences[0]:
for y in icross(*sequences[1:]):
yield (x,)+y
else: yield ()
[docs]def get_groups(iterable, n, step):
"""Make groups of 'n' elements from the iterable advancing
'step' elements each iteration"""
itlist = tee(iterable, n)
onestepit = izip(*(starmap(drop, enumerate(itlist))))
return takeevery(step, onestepit)
[docs]def quantify(iterable, pred=bool):
"Count how many times the predicate is true"
return sum(imap(pred, iterable))
[docs]def pairwise(iterable):
"s -> (s0,s1), (s1,s2), (s2, s3), ..."
a, b = tee(iterable)
next(b, None)
return izip(a, b)
[docs]def rand_seq(size):
'''generates values in random order
equivalent to using shuffle in random,
without generating all values at once'''
values=range(size)
for i in xrange(size):
# pick a random index into remaining values
j=i+int(random.random()*(size-i))
# swap the values
values[j],values[i]=values[i],values[j]
# return the swapped value
yield values[i]
[docs]def all_pairs(size):
'''generates all i,j pairs for i,j from 0-size'''
for i in rand_seq(size):
for j in rand_seq(size):
yield (i,j)
[docs]def split(iterable,condition):
"""@return list of elements in iterable that satisfy condition, and those that don't"""
yes,no=[],[]
for x in iterable:
if condition(x):
yes.append(x)
else:
no.append(x)
return yes,no
[docs]def next_permutation(seq, pred=cmp):
"""Like C++ std::next_permutation() but implemented as
generator. Yields copies of seq.
see http://blog.bjrn.se/2008/04/lexicographic-permutations-using.html"""
def reverse(seq, start, end):
# seq = seq[:start] + reversed(seq[start:end]) + \
# seq[end:]
end -= 1
if end <= start:
return
while True:
seq[start], seq[end] = seq[end], seq[start]
if start == end or start+1 == end:
return
start += 1
end -= 1
if not seq:
raise StopIteration
try:
seq[0]
except TypeError:
raise TypeError("seq must allow random access.")
first = 0
last = len(seq)
seq = seq[:]
# Yield input sequence as the STL version is often
# used inside do {} while.
yield seq
if last == 1:
raise StopIteration
while True:
next = last - 1
while True:
# Step 1.
next1 = next
next -= 1
if pred(seq[next], seq[next1]) < 0:
# Step 2.
mid = last - 1
while not (pred(seq[next], seq[mid]) < 0):
mid -= 1
seq[next], seq[mid] = seq[mid], seq[next]
# Step 3.
reverse(seq, next1, last)
# Change to yield references to get rid of
# (at worst) |seq|! copy operations.
yield seq[:]
break
if next == first:
raise StopIteration
raise StopIteration
[docs]class iter2(object):
"""Takes in an object that is iterable.
http://code.activestate.com/recipes/578092-flattening-an-arbitrarily-deep-list-or-any-iterato/
Allows for the following method calls (that should be built into iterators anyway...)
calls:
- append - appends another iterable onto the iterator.
- insert - only accepts inserting at the 0 place, inserts an iterable before other iterables.
- adding. an iter2 object can be added to another object that is
iterable. i.e. iter2 + iter (not iter + iter2).
It's best to make all objects iter2 objects to avoid syntax errors. :D
"""
[docs] def __init__(self, iterable):
self._iter = iter(iterable)
[docs] def append(self, iterable):
self._iter = chain(self._iter, iter(iterable))
[docs] def insert(self, place, iterable):
if place != 0:
raise ValueError('Can only insert at index of 0')
self._iter = chain(iter(iterable), self._iter)
[docs] def __add__(self, iterable):
return chain(self._iter, iter(iterable))
[docs] def next(self):
return self._iter.next()
[docs] def __iter__(self):
return self
[docs]def iflatten(iterable):
'''flatten a list of any depth'''
iterable = iter2(iterable)
for e in iterable:
if hasattr(e, '__iter__'):
iterable.insert(0, e)
else:
yield e