iterables.py

import collections
import itertools
import operator

from ..core.compatibility import as_int
from .enumerative import (MultisetPartitionTraverser, list_visitor,
                          multiset_partitions_taocp)


def flatten(iterable, levels=None, cls=None):
    """
    Recursively denest iterable containers.

    >>> flatten([1, 2, 3])
    [1, 2, 3]
    >>> flatten([1, 2, [3]])
    [1, 2, 3]
    >>> flatten([1, [2, 3], [4, 5]])
    [1, 2, 3, 4, 5]
    >>> flatten([1.0, 2, (1, None)])
    [1.0, 2, 1, None]

    If you want to denest only a specified number of levels of
    nested containers, then set ``levels`` flag to the desired
    number of levels::

    >>> ls = [[(-2, -1), (1, 2)], [(0, 0)]]

    >>> flatten(ls, levels=1)
    [(-2, -1), (1, 2), (0, 0)]

    If cls argument is specified, it will only flatten instances of that
    class, for example:

    >>> class MyOp(Basic):
    ...     pass
    ...
    >>> flatten([MyOp(1, MyOp(2, 3))], cls=MyOp)
    [1, 2, 3]

    adapted from https://kogs-www.informatik.uni-hamburg.de/~meine/python_tricks

    """
    if levels is not None:
        if not levels:
            return iterable
        if levels > 0:
            levels -= 1
        else:
            raise ValueError(
                f'expected non-negative number of levels, got {levels}')

    if cls is None:
        def reducible(x):
            return is_sequence(x, set)
    else:
        def reducible(x):
            return isinstance(x, cls)

    result = []

    for el in iterable:
        if reducible(el):
            if hasattr(el, 'args'):
                el = el.args
            result.extend(flatten(el, levels=levels, cls=cls))
        else:
            result.append(el)

    return result


def unflatten(iter, n=2):
    """Group ``iter`` into tuples of length ``n``. Raise an error if
    the length of ``iter`` is not a multiple of ``n``.

    """
    if n < 1 or len(iter) % n:
        raise ValueError(f'iter length is not a multiple of {n:d}')
    return list(zip(*(iter[i::n] for i in range(n))))


def group(seq, multiple=True):
    """
    Splits a sequence into a list of lists of equal, adjacent elements.

    Examples
    ========

    >>> group([1, 1, 1, 2, 2, 3])
    [[1, 1, 1], [2, 2], [3]]
    >>> group([1, 1, 1, 2, 2, 3], multiple=False)
    [(1, 3), (2, 2), (3, 1)]
    >>> group([1, 1, 3, 2, 2, 1], multiple=False)
    [(1, 2), (3, 1), (2, 2), (1, 1)]

    See Also
    ========
    multiset

    """
    if not seq:
        return []

    current, groups = [seq[0]], []

    for elem in seq[1:]:
        if elem == current[-1]:
            current.append(elem)
        else:
            groups.append(current)
            current = [elem]

    groups.append(current)

    if multiple:
        return groups

    for i, current in enumerate(groups):
        groups[i] = (current[0], len(current))

    return groups


def multiset(seq):
    """Return the hashable sequence in multiset form with values being the
    multiplicity of the item in the sequence.

    Examples
    ========

    >>> multiset('mississippi')
    {'i': 4, 'm': 1, 'p': 2, 's': 4}

    See Also
    ========

    group

    """
    rv = collections.defaultdict(int)
    for s in seq:
        rv[s] += 1
    return dict(rv)


def postorder_traversal(node, keys=None):
    """
    Do a postorder traversal of a tree.

    This generator recursively yields nodes that it has visited in a postorder
    fashion. That is, it descends through the tree depth-first to yield all of
    a node's children's postorder traversal before yielding the node itself.

    Parameters
    ==========

    node : diofant expression
        The expression to traverse.
    keys : (default None) sort key(s)
        The key(s) used to sort args of Basic objects. When None, args of Basic
        objects are processed in arbitrary order. If key is defined, it will
        be passed along to ordered() as the only key(s) to use to sort the
        arguments; if ``key`` is simply True then the default keys of
        ``ordered`` will be used (node count and default_sort_key).

    Yields
    ======
    subtree : diofant expression
        All of the subtrees in the tree.

    Examples
    ========

    >>> from diofant.abc import w

    The nodes are returned in the order that they are encountered unless key
    is given; simply passing key=True will guarantee that the traversal is
    unique.

    >>> list(postorder_traversal(w + (x + y)*z, keys=True))
    [w, z, x, y, x + y, z*(x + y), w + z*(x + y)]

    """
    from ..core import Basic

    if isinstance(node, Basic):
        args = node.args
        if keys:
            if keys != 1:
                args = ordered(args, keys, default=False)
            else:
                args = ordered(args)
        for arg in args:
            yield from postorder_traversal(arg, keys)
    elif is_iterable(node):
        for item in node:
            yield from postorder_traversal(item, keys)
    yield node


def subsets(seq, k=None, repetition=False):
    """Generates all k-subsets (combinations) from an n-element set, seq.

    A k-subset of an n-element set is any subset of length exactly k. The
    number of k-subsets of an n-element set is given by binomial(n, k),
    whereas there are 2**n subsets all together. If k is None then all
    2**n subsets will be returned from shortest to longest.

    Examples
    ========

    subsets(seq, k) will return the n!/k!/(n - k)! k-subsets (combinations)
    without repetition, i.e. once an item has been removed, it can no
    longer be "taken":

    >>> from diofant.utilities.iterables import subsets

    >>> list(subsets([1, 2], 2))
    [(1, 2)]
    >>> list(subsets([1, 2]))
    [(), (1,), (2,), (1, 2)]
    >>> list(subsets([1, 2, 3], 2))
    [(1, 2), (1, 3), (2, 3)]

    subsets(seq, k, repetition=True) will return the (n - 1 + k)!/k!/(n - 1)!
    combinations *with* repetition:

    >>> list(subsets([1, 2], 2, repetition=True))
    [(1, 1), (1, 2), (2, 2)]

    If you ask for more items than are in the set you get the empty set unless
    you allow repetitions:

    >>> list(subsets([0, 1], 3, repetition=False))
    []
    >>> list(subsets([0, 1], 3, repetition=True))
    [(0, 0, 0), (0, 0, 1), (0, 1, 1), (1, 1, 1)]

    """
    if not repetition:
        meth = itertools.combinations
    else:
        meth = itertools.combinations_with_replacement
    if k is None:
        yield from itertools.chain.from_iterable(meth(seq, i)
                                                 for i in range(len(seq) + 1))
    else:
        yield from meth(seq, k)


def numbered_symbols(prefix='x', cls=None, start=0, exclude=[], **assumptions):
    """
    Generate an infinite stream of Symbols consisting of a prefix and
    increasing subscripts provided that they do not occur in `exclude`.

    Parameters
    ==========

    prefix : str, optional
        The prefix to use. By default, this function will generate symbols of
        the form "x0", "x1", etc.

    cls : class, optional
        The class to use. By default, it uses Symbol, but you can also use Wild or Dummy.

    start : int, optional
        The start number.  By default, it is 0.

    Returns
    =======

    sym : Symbol
        The subscripted symbols.

    """
    exclude = set(exclude or [])
    if cls is None:
        # We can't just make the default cls=Symbol because it isn't
        # imported yet.
        from ..core import Symbol
        cls = Symbol

    while True:
        name = f'{prefix}{start}'
        s = cls(name, **assumptions)
        if s not in exclude:
            yield s
        start += 1


def sift(seq, keyfunc):
    """
    Sift the sequence, ``seq`` into a dictionary according to keyfunc.

    OUTPUT: each element in expr is stored in a list keyed to the value
    of keyfunc for the element.

    Examples
    ========

    >>> from collections import defaultdict

    >>> sift(range(5), lambda x: x % 2) == defaultdict(int, {0: [0, 2, 4], 1: [1, 3]})
    True

    sift() returns a defaultdict() object, so any key that has no matches will
    give [].

    >>> dl = sift([x], lambda x: x.is_commutative)
    >>> dl == defaultdict(list, {True: [x]})
    True
    >>> dl[False]
    []

    Sometimes you won't know how many keys you will get:

    >>> (sift([sqrt(x), exp(x), (y**x)**2],
    ...       lambda x: x.as_base_exp()[0]) ==
    ...  defaultdict(list, {E: [exp(x)], x: [sqrt(x)], y: [y**(2*x)]}))
    True

    If you need to sort the sifted items it might be better to use
    ``ordered`` which can economically apply multiple sort keys
    to a squence while sorting.

    See Also
    ========

    ordered

    """
    m = collections.defaultdict(list)
    for i in seq:
        m[keyfunc(i)].append(i)
    return m


def common_prefix(*seqs):
    """Return the subsequence that is a common start of sequences in ``seqs``.

    >>> common_prefix(list(range(3)))
    [0, 1, 2]
    >>> common_prefix(list(range(3)), list(range(4)))
    [0, 1, 2]
    >>> common_prefix([1, 2, 3], [1, 2, 5])
    [1, 2]
    >>> common_prefix([1, 2, 3], [1, 3, 5])
    [1]

    """
    if any(not s for s in seqs):
        return []
    if len(seqs) == 1:
        return seqs[0]
    i = 0
    for i in range(min(len(s) for s in seqs)):
        if not all(seqs[j][i] == seqs[0][i] for j in range(len(seqs))):
            break
    else:
        i += 1
    return seqs[0][:i]


def common_suffix(*seqs):
    """Return the subsequence that is a common ending of sequences in ``seqs``.

    >>> common_suffix(list(range(3)))
    [0, 1, 2]
    >>> common_suffix(list(range(3)), list(range(4)))
    []
    >>> common_suffix([1, 2, 3], [9, 2, 3])
    [2, 3]
    >>> common_suffix([1, 2, 3], [9, 7, 3])
    [3]

    """
    if any(not s for s in seqs):
        return []
    if len(seqs) == 1:
        return seqs[0]
    i = 0
    for i in range(-1, -min(len(s) for s in seqs) - 1, -1):
        if not all(seqs[j][i] == seqs[0][i] for j in range(len(seqs))):
            break
    else:
        i -= 1
    if i == -1:
        return []
    return seqs[0][i + 1:]


def rotate_left(x, y):
    """
    Left rotates a list x by the number of steps specified
    in y.

    Examples
    ========

    >>> a = [0, 1, 2]
    >>> rotate_left(a, 1)
    [1, 2, 0]

    """
    if len(x) == 0:
        return []
    y = y % len(x)
    return x[y:] + x[:y]


def rotate_right(x, y):
    """
    Right rotates a list x by the number of steps specified
    in y.

    Examples
    ========

    >>> a = [0, 1, 2]
    >>> rotate_right(a, 1)
    [2, 0, 1]

    """
    if len(x) == 0:
        return []
    y = len(x) - y % len(x)
    return x[y:] + x[:y]


def multiset_combinations(m, n, g=None):
    """
    Return the unique combinations of size ``n`` from multiset ``m``.

    Examples
    ========

    >>> from itertools import combinations
    >>> [''.join(i) for i in multiset_combinations('baby', 3)]
    ['abb', 'aby', 'bby']

    >>> def count(f, s):
    ...     return len(list(f(s, 3)))

    The number of combinations depends on the number of letters; the
    number of unique combinations depends on how the letters are
    repeated.

    >>> s1 = 'abracadabra'
    >>> s2 = 'banana tree'
    >>> count(combinations, s1), count(multiset_combinations, s1)
    (165, 23)
    >>> count(combinations, s2), count(multiset_combinations, s2)
    (165, 54)

    """
    if g is None:
        if type(m) is dict:
            if n > sum(m.values()):
                return
            g = [[k, m[k]] for k in ordered(m)]
        else:
            m = list(m)
            if n > len(m):
                return
            try:
                m = multiset(m)
                g = [(k, m[k]) for k in ordered(m)]
            except TypeError:
                m = list(ordered(m))
                g = [list(i) for i in group(m, multiple=False)]
        del m
    if sum(v for k, v in g) < n or not n:
        yield []
    else:
        for i, (k, v) in enumerate(g):
            if v >= n:
                yield [k]*n
                v = n - 1
            for v in range(min(n, v), 0, -1):
                for j in multiset_combinations(None, n - v, g[i + 1:]):
                    rv = [k]*v + j
                    if len(rv) == n:
                        yield rv


def multiset_permutations(m, size=None, g=None):
    """
    Return the unique permutations of multiset ``m``.

    Examples
    ========

    >>> [''.join(i) for i in multiset_permutations('aab')]
    ['aab', 'aba', 'baa']
    >>> factorial(len('banana'))
    720
    >>> len(list(multiset_permutations('banana')))
    60

    """
    if g is None:
        if type(m) is dict:
            g = [[k, m[k]] for k in ordered(m)]
        else:
            m = list(ordered(m))
            g = [list(i) for i in group(m, multiple=False)]
        del m
    do = [gi for gi in g if gi[1] > 0]
    SUM = sum(gi[1] for gi in do)
    if not do or size is not None and (size > SUM or size < 1):
        if size < 1:
            yield []
    elif size == 1:
        for k, v in do:
            yield [k]
    elif len(do) == 1:
        k, v = do[0]
        v = v if size is None else (size if size <= v else 0)
        yield [k for i in range(v)]
    elif all(v == 1 for k, v in do):
        for p in itertools.permutations([k for k, v in do], size):
            yield list(p)
    else:
        size = size if size is not None else SUM
        for i, (k, v) in enumerate(do):
            do[i][1] -= 1
            for j in multiset_permutations(None, size - 1, do):
                yield [k] + j
            do[i][1] += 1


def _partition(seq, vector, m=None):
    """
    Return the partion of seq as specified by the partition vector.

    Examples
    ========

    >>> _partition('abcde', [1, 0, 1, 2, 0])
    [['b', 'e'], ['a', 'c'], ['d']]

    Specifying the number of bins in the partition is optional:

    >>> _partition('abcde', [1, 0, 1, 2, 0], 3)
    [['b', 'e'], ['a', 'c'], ['d']]

    The output of _set_partitions can be passed as follows:

    >>> output = (3, [1, 0, 1, 2, 0])
    >>> _partition('abcde', *output)
    [['b', 'e'], ['a', 'c'], ['d']]

    See Also
    ========
    combinatorics.partitions.Partition.from_rgs()

    """
    if m is None:
        m = max(vector) + 1
    elif type(vector) is int:  # entered as m, vector
        vector, m = m, vector
    p = [[] for i in range(m)]
    for i, v in enumerate(vector):
        p[v].append(seq[i])
    return p


def _set_partitions(n):
    """Cycle through all partitions of n elements, yielding the
    current number of partitions, ``m``, and a mutable list, ``q``
    such that element[i] is in part q[i] of the partition.

    NOTE: ``q`` is modified in place and generally should not be changed
    between function calls.

    Examples
    ========

    >>> for m, q in _set_partitions(3):
    ...     print(f'{m} {q} { _partition("abc", q, m)}')
    1 [0, 0, 0] [['a', 'b', 'c']]
    2 [0, 0, 1] [['a', 'b'], ['c']]
    2 [0, 1, 0] [['a', 'c'], ['b']]
    2 [0, 1, 1] [['a'], ['b', 'c']]
    3 [0, 1, 2] [['a'], ['b'], ['c']]

    Notes
    =====

    This algorithm is similar to, and solves the same problem as,
    Algorithm 7.2.1.5H, from volume 4A of Knuth's The Art of Computer
    Programming.  Knuth uses the term "restricted growth string" where
    this code refers to a "partition vector". In each case, the meaning is
    the same: the value in the ith element of the vector specifies to
    which part the ith set element is to be assigned.

    At the lowest level, this code implements an n-digit big-endian
    counter (stored in the array q) which is incremented (with carries) to
    get the next partition in the sequence.  A special twist is that a
    digit is constrained to be at most one greater than the maximum of all
    the digits to the left of it.  The array p maintains this maximum, so
    that the code can efficiently decide when a digit can be incremented
    in place or whether it needs to be reset to 0 and trigger a carry to
    the next digit.  The enumeration starts with all the digits 0 (which
    corresponds to all the set elements being assigned to the same 0th
    part), and ends with 0123...n, which corresponds to each set element
    being assigned to a different, singleton, part.

    This routine was rewritten to use 0-based lists while trying to
    preserve the beauty and efficiency of the original algorithm.

    References
    ==========

    * Nijenhuis, Albert and Wilf, Herbert. (1978) Combinatorial Algorithms,
      2nd Ed, p 91, algorithm "nexequ". Available online from
      http://www.math.upenn.edu/~wilf/website/CombAlgDownld.html (viewed
      November 17, 2012).

    """
    p = [0]*n
    q = [0]*n
    nc = 1
    yield nc, q
    while nc != n:
        m = n
        while 1:
            m -= 1
            i = q[m]
            if p[i] != 1:
                break
            q[m] = 0
        i += 1
        q[m] = i
        m += 1
        nc += m - n
        p[0] += n - m
        if i == nc:
            p[nc] = 0
            nc += 1
        p[i - 1] -= 1
        p[i] += 1
        yield nc, q


def multiset_partitions(multiset, m=None):
    """
    Return unique partitions of the given multiset (in list form).
    If ``m`` is None, all multisets will be returned, otherwise only
    partitions with ``m`` parts will be returned.

    If ``multiset`` is an integer, a range [0, 1, ..., multiset - 1]
    will be supplied.

    *Counting*

    The number of partitions of a set is given by the bell number:

    >>> len(list(multiset_partitions(5))) == bell(5) == 52
    True

    The number of partitions of length k from a set of size n is given by the
    Stirling Number of the 2nd kind:

    >>> def s2(n, k):
    ...     from diofant import Dummy, Sum, binomial, factorial
    ...     if k > n:
    ...         return 0
    ...     j = Dummy()
    ...     arg = (-1)**(k-j)*j**n*binomial(k, j)
    ...     return 1/factorial(k)*Sum(arg, (j, 0, k)).doit()
    ...
    >>> s2(5, 2) == len(list(multiset_partitions(5, 2))) == 15
    True

    These comments on counting apply to *sets*, not multisets.

    Examples
    ========

    >>> list(multiset_partitions([1, 2, 3, 4], 2))
    [[[1, 2, 3], [4]], [[1, 2, 4], [3]], [[1, 2], [3, 4]],
    [[1, 3, 4], [2]], [[1, 3], [2, 4]], [[1, 4], [2, 3]],
    [[1], [2, 3, 4]]]
    >>> list(multiset_partitions([1, 2, 3, 4], 1))
    [[[1, 2, 3, 4]]]

    Only unique partitions are returned and these will be returned in a
    canonical order regardless of the order of the input:

    >>> a = [1, 2, 2, 1]
    >>> ans = list(multiset_partitions(a, 2))
    >>> a.sort()
    >>> list(multiset_partitions(a, 2)) == ans
    True
    >>> a = range(3, 1, -1)
    >>> (list(multiset_partitions(a)) ==
    ...  list(multiset_partitions(sorted(a))))
    True

    If m is omitted then all partitions will be returned:

    >>> list(multiset_partitions([1, 1, 2]))
    [[[1, 1, 2]], [[1, 1], [2]], [[1, 2], [1]], [[1], [1], [2]]]
    >>> list(multiset_partitions([1]*3))
    [[[1, 1, 1]], [[1], [1, 1]], [[1], [1], [1]]]

    Notes
    =====

    When all the elements are the same in the multiset, the order
    of the returned partitions is determined by the ``partitions``
    routine. If one is counting partitions then it is better to use
    the ``nT`` function.

    See Also
    ========

    partitions
    diofant.combinatorics.partitions.Partition
    diofant.combinatorics.partitions.IntegerPartition

    """
    # This function looks at the supplied input and dispatches to
    # several special-case routines as they apply.
    if type(multiset) is int:
        n = multiset
        if m and m > n:
            return
        multiset = list(range(n))
        if m == 1:
            yield [multiset[:]]
            return

        # If m is not None, it can sometimes be faster to use
        # MultisetPartitionTraverser.enum_range() even for inputs
        # which are sets.  Since the _set_partitions code is quite
        # fast, this is only advantageous when the overall set
        # partitions outnumber those with the desired number of parts
        # by a large factor.  (At least 60.)  Such a switch is not
        # currently implemented.
        for nc, q in _set_partitions(n):
            if m is None or nc == m:
                rv = [[] for i in range(nc)]
                for i in range(n):
                    rv[q[i]].append(multiset[i])
                yield rv
        return

    if len(multiset) == 1 and type(multiset) is str:
        multiset = [multiset]

    if len(set(multiset)) == 1:
        # Only one component, repeated n times.  The resulting
        # partitions correspond to partitions of integer n.
        n = len(multiset)
        if m and m > n:
            return
        if m == 1:
            yield [multiset[:]]
            return
        x = multiset[:1]
        for size, p in partitions(n, m, size=True):
            if m is None or size == m:
                rv = []
                for k in sorted(p):
                    rv.extend([x*k]*p[k])
                yield rv
    else:
        multiset = list(ordered(multiset))
        n = len(multiset)
        if m and m > n:
            return
        if m == 1:
            yield [multiset[:]]
            return

        # Split the information of the multiset into two lists -
        # one of the elements themselves, and one (of the same length)
        # giving the number of repeats for the corresponding element.
        elements, multiplicities = zip(*group(multiset, False))

        if len(elements) < len(multiset):
            # General case - multiset with more than one distinct element
            # and at least one element repeated more than once.
            if m:
                mpt = MultisetPartitionTraverser()
                for state in mpt.enum_range(multiplicities, m-1, m):
                    yield list_visitor(state, elements)
            else:
                for state in multiset_partitions_taocp(multiplicities):
                    yield list_visitor(state, elements)
        else:
            # Set partitions case - no repeated elements. Pretty much
            # same as int argument case above, with same possible, but
            # currently unimplemented optimization for some cases when
            # m is not None
            for nc, q in _set_partitions(n):
                if m is None or nc == m:
                    rv = [[] for i in range(nc)]
                    for i in range(n):
                        rv[q[i]].append(i)
                    yield [[multiset[j] for j in i] for i in rv]


def partitions(n, m=None, k=None, size=False):
    """Generate all partitions of positive integer, n.

    Parameters
    ==========

    ``m`` : integer (default gives partitions of all sizes)
        limits number of parts in partition (mnemonic: m, maximum parts).
        Default value, None, gives partitions from 1 through n.
    ``k`` : integer (default gives partitions number from 1 through n)
        limits the numbers that are kept in the partition (mnemonic: k, keys)
    ``size`` : bool (default False, only partition is returned)
        when ``True`` then (M, P) is returned where M is the sum of the
        multiplicities and P is the generated partition.

    Each partition is represented as a dictionary, mapping an integer
    to the number of copies of that integer in the partition.  For example,
    the first partition of 4 returned is {4: 1}, "4: one of them".

    Examples
    ========

    The numbers appearing in the partition (the key of the returned dict)
    are limited with k:

    >>> from diofant.utilities.iterables import partitions

    >>> for p in partitions(6, k=2):
    ...     print(p)
    {2: 3}
    {2: 2, 1: 2}
    {2: 1, 1: 4}
    {1: 6}

    The maximum number of parts in the partition (the sum of the values in
    the returned dict) are limited with m:

    >>> for p in partitions(6, m=2):
    ...     print(p)
    ...
    {6: 1}
    {5: 1, 1: 1}
    {4: 1, 2: 1}
    {3: 2}

    Note that the _same_ dictionary object is returned each time.
    This is for speed:  generating each partition goes quickly,
    taking constant time, independent of n.

    >>> list(partitions(6, k=2))
    [{1: 6}, {1: 6}, {1: 6}, {1: 6}]

    If you want to build a list of the returned dictionaries then
    make a copy of them:

    >>> [p.copy() for p in partitions(6, k=2)]
    [{2: 3}, {1: 2, 2: 2}, {1: 4, 2: 1}, {1: 6}]
    >>> [(M, p.copy()) for M, p in partitions(6, k=2, size=True)]
    [(3, {2: 3}), (4, {1: 2, 2: 2}), (5, {1: 4, 2: 1}), (6, {1: 6})]

    References
    ==========

    * http://code.activestate.com/recipes/218332-generator-for-integer-partitions/

    Notes
    =====

    Modified from Tim Peter's version to allow for k and m values.

    See Also
    ========

    diofant.combinatorics.partitions.Partition
    diofant.combinatorics.partitions.IntegerPartition

    """
    if (n <= 0 or m is not None and m < 1 or
            k is not None and k < 1 or m and k and m*k < n):
        # the empty set is the only way to handle these inputs
        if size:
            yield 0, {}
        else:
            yield {}
        return

    if m is None:
        m = n
    else:
        m = min(m, n)

    k = min(k or n, n)

    n, m, k = as_int(n), as_int(m), as_int(k)
    q, r = divmod(n, k)
    ms = {k: q}
    keys = [k]  # ms keys, from largest to smallest
    if r:
        ms[r] = 1
        keys.append(r)
    room = m - q - bool(r)
    if size:
        yield sum(ms.values()), ms
    else:
        yield ms

    while keys != [1]:
        # Reuse any 1's.
        if keys[-1] == 1:
            del keys[-1]
            reuse = ms.pop(1)
            room += reuse
        else:
            reuse = 0

        while 1:
            # Let i be the smallest key larger than 1.  Reuse one
            # instance of i.
            i = keys[-1]
            newcount = ms[i] = ms[i] - 1
            reuse += i
            if newcount == 0:
                del keys[-1], ms[i]
            room += 1

            # Break the remainder into pieces of size i-1.
            i -= 1
            q, r = divmod(reuse, i)
            need = q + bool(r)
            if need > room:
                if not keys:
                    return
                continue

            ms[i] = q
            keys.append(i)
            if r:
                ms[r] = 1
                keys.append(r)
            break
        room -= need
        if size:
            yield sum(ms.values()), ms
        else:
            yield ms


def ordered_partitions(n, m=None, sort=True):
    """Generates ordered partitions of integer ``n``.

    Parameters
    ==========

    ``m`` : int or None, optional
        By default (None) gives partitions of all sizes, else only
        those with size m. In addition, if ``m`` is not None then
        partitions are generated *in place* (see examples).
    ``sort`` : bool, optional
        Controls whether partitions are
        returned in sorted order (default) when ``m`` is not None; when False,
        the partitions are returned as fast as possible with elements
        sorted, but when m|n the partitions will not be in
        ascending lexicographical order.

    Examples
    ========

    All partitions of 5 in ascending lexicographical:

    >>> for p in ordered_partitions(5):
    ...     print(p)
    [1, 1, 1, 1, 1]
    [1, 1, 1, 2]
    [1, 1, 3]
    [1, 2, 2]
    [1, 4]
    [2, 3]
    [5]

    Only partitions of 5 with two parts:

    >>> for p in ordered_partitions(5, 2):
    ...     print(p)
    [1, 4]
    [2, 3]

    When ``m`` is given, a given list objects will be used more than
    once for speed reasons so you will not see the correct partitions
    unless you make a copy of each as it is generated:

    >>> list(ordered_partitions(7, 3))
    [[1, 1, 1], [1, 1, 1], [1, 1, 1], [2, 2, 2]]
    >>> [list(p) for p in ordered_partitions(7, 3)]
    [[1, 1, 5], [1, 2, 4], [1, 3, 3], [2, 2, 3]]

    When ``n`` is a multiple of ``m``, the elements are still sorted
    but the partitions themselves will be *unordered* if sort is False;
    the default is to return them in ascending lexicographical order.

    >>> for p in ordered_partitions(6, 2):
    ...     print(p)
    [1, 5]
    [2, 4]
    [3, 3]

    But if speed is more important than ordering, sort can be set to
    False:

    >>> for p in ordered_partitions(6, 2, sort=False):
    ...     print(p)
    [1, 5]
    [3, 3]
    [2, 4]

    References
    ==========

    * Generating Integer Partitions, [online],
      Available: http://jeromekelleher.net/generating-integer-partitions.html
    * Jerome Kelleher and Barry O'Sullivan, "Generating All
      Partitions: A Comparison Of Two Encodings", [online],
      Available: https://arxiv.org/pdf/0909.2331v2.pdf

    """
    if n < 1 or m is not None and m < 1:
        # the empty set is the only way to handle these inputs
        yield []
        return

    if m is None:
        # The list `a`'s leading elements contain the partition in which
        # y is the biggest element and x is either the same as y or the
        # 2nd largest element; v and w are adjacent element indices
        # to which x and y are being assigned, respectively.
        a = [1]*n
        y = -1
        v = n
        while v > 0:
            v -= 1
            x = a[v] + 1
            while y >= 2 * x:
                a[v] = x
                y -= x
                v += 1
            w = v + 1
            while x <= y:
                a[v] = x
                a[w] = y
                yield a[:w + 1]
                x += 1
                y -= 1
            a[v] = x + y
            y = a[v] - 1
            yield a[:w]
    elif m == 1:
        yield [n]
    elif n == m:
        yield [1]*n
    else:
        # recursively generate partitions of size m
        for b in range(1, n//m + 1):
            a = [b]*m
            x = n - b*m
            if not x:
                if sort:
                    yield a
            elif not sort and x <= m:
                for ax in ordered_partitions(x, sort=False):
                    mi = len(ax)
                    a[-mi:] = [i + b for i in ax]
                    yield a
                    a[-mi:] = [b]*mi
            else:
                for mi in range(1, m):
                    for ax in ordered_partitions(x, mi, sort=True):
                        a[-mi:] = [i + b for i in ax]
                        yield a
                        a[-mi:] = [b]*mi


def uniq(seq, result=None):
    """
    Yield unique elements from ``seq`` as an iterator. The second
    parameter ``result``  is used internally; it is not necessary to pass
    anything for this.

    Examples
    ========

    >>> dat = [1, 4, 1, 5, 4, 2, 1, 2]
    >>> type(uniq(dat)) in (list, tuple)
    False

    >>> list(uniq(dat))
    [1, 4, 5, 2]
    >>> list(uniq(x for x in dat))
    [1, 4, 5, 2]
    >>> list(uniq([[1], [2, 1], [1]]))
    [[1], [2, 1]]

    """
    try:
        seen = set()
        result = result or []
        for i, s in enumerate(seq):
            if not (s in seen or seen.add(s)):
                yield s
    except TypeError:
        if s not in result:
            yield s
            result.append(s)
        if hasattr(seq, '__getitem__'):
            yield from uniq(seq[i + 1:], result)  # pylint: disable=used-before-assignment
        else:
            yield from uniq(seq, result)


def minlex(seq, directed=True, is_set=False, small=None):
    """
    Return a tuple where the smallest element appears first; if
    ``directed`` is True (default) then the order is preserved, otherwise
    the sequence will be reversed if that gives a smaller ordering.

    If every element appears only once then is_set can be set to True
    for more efficient processing.

    If the smallest element is known at the time of calling, it can be
    passed and the calculation of the smallest element will be omitted.

    Examples
    ========

    >>> minlex((1, 2, 0))
    (0, 1, 2)
    >>> minlex((1, 0, 2))
    (0, 2, 1)
    >>> minlex((1, 0, 2), directed=False)
    (0, 1, 2)

    >>> minlex('11010011000', directed=True)
    '00011010011'
    >>> minlex('11010011000', directed=False)
    '00011001011'

    """
    is_str = isinstance(seq, str)
    seq = list(seq)
    if small is None:
        small = min(seq, key=default_sort_key)
    if is_set:
        i = seq.index(small)
        if not directed:
            n = len(seq)
            p = (i + 1) % n
            m = (i - 1) % n
            if default_sort_key(seq[p]) > default_sort_key(seq[m]):
                seq = list(reversed(seq))
                i = n - i - 1
        if i:
            seq = rotate_left(seq, i)
        best = seq
    else:
        count = seq.count(small)
        if count == 1 and directed:
            best = rotate_left(seq, seq.index(small))
        else:
            # if not directed, and not a set, we can't just
            # pass this off to minlex with is_set True since
            # peeking at the neighbor may not be sufficient to
            # make the decision so we continue...
            best = seq
            for i in range(count):
                seq = rotate_left(seq, seq.index(small, count != 1))
                best = min(best, seq)
                # it's cheaper to rotate now rather than search
                # again for these in reversed order so we test
                # the reverse now
                if not directed:
                    seq = rotate_left(seq, 1)
                    seq = list(reversed(seq))
                    best = min(best, seq)
                    seq = list(reversed(seq))
                    seq = rotate_right(seq, 1)
    # common return
    if is_str:
        return ''.join(best)
    return tuple(best)


def runs(seq, op=operator.gt):
    """Group the sequence into lists in which successive elements
    all compare the same with the comparison operator, ``op``:
    op(seq[i + 1], seq[i]) is True from all elements in a run.

    Examples
    ========

    >>> import operator
    >>> runs([0, 1, 2, 2, 1, 4, 3, 2, 2])
    [[0, 1, 2], [2], [1, 4], [3], [2], [2]]
    >>> runs([0, 1, 2, 2, 1, 4, 3, 2, 2], op=operator.ge)
    [[0, 1, 2, 2], [1, 4], [3], [2, 2]]

    """
    cycles = []
    seq = iter(seq)
    try:
        run = [next(seq)]
    except StopIteration:
        return []
    while True:
        try:
            ei = next(seq)
        except StopIteration:
            break
        if op(ei, run[-1]):
            run.append(ei)
            continue
        cycles.append(run)
        run = [ei]
    cycles.append(run)
    return cycles


def cantor_product(*args):
    """Breadth-first (diagonal) cartesian product of iterables.

    Each iterable is advanced in turn in a round-robin fashion. As usual with
    breadth-first, this comes at the cost of memory consumption.

    >>> from itertools import count, islice
    >>> list(islice(cantor_product(count(), count()), 9))
    [(0, 0), (0, 1), (1, 0), (1, 1), (0, 2), (1, 2), (2, 0), (2, 1), (2, 2)]

    """
    args = list(map(iter, args))

    try:
        argslist = [[next(a)] for a in args]
    except StopIteration:
        return
    yield tuple(a[0] for a in argslist)

    nargs = len(args)
    exhausted = [False]*nargs
    n = nargs

    while not all(exhausted):
        n = (n - 1) % nargs
        if not exhausted[n]:
            try:
                argslist[n].append(next(args[n]))
            except StopIteration:
                exhausted[n] = True
            else:
                yield from itertools.product(*(argslist[:n] + [argslist[n][-1:]] +
                                               argslist[n + 1:]))


def permute_signs(t):
    """Return iterator in which the signs of non-zero elements
    of t are permuted.

    Examples
    ========

    >>> list(permute_signs((0, 1, 2)))
    [(0, 1, 2), (0, -1, 2), (0, 1, -2), (0, -1, -2)]

    """
    for signs in itertools.product(*[(1, -1)]*(len(t) - t.count(0))):
        signs = list(signs)
        yield type(t)([i*signs.pop() if i else i for i in t])


def signed_permutations(t):
    """Return iterator in which the signs of non-zero elements
    of t and the order of the elements are permuted.

    Examples
    ========

    >>> list(signed_permutations((0, 1, 2)))
    [(0, 1, 2), (0, -1, 2), (0, 1, -2), (0, -1, -2), (0, 2, 1),
     (0, -2, 1), (0, 2, -1), (0, -2, -1), (1, 0, 2), (-1, 0, 2),
     (1, 0, -2), (-1, 0, -2), (1, 2, 0), (-1, 2, 0), (1, -2, 0),
     (-1, -2, 0), (2, 0, 1), (-2, 0, 1), (2, 0, -1), (-2, 0, -1),
     (2, 1, 0), (-2, 1, 0), (2, -1, 0), (-2, -1, 0)]

    """
    return (type(t)(i) for j in itertools.permutations(t)
            for i in permute_signs(j))


def _nodes(e):
    """
    A helper for ordered() which returns the node count of ``e`` which
    for Basic objects is the number of Basic nodes in the expression tree
    but for other objects is 1 (unless the object is an iterable or dict
    for which the sum of nodes is returned).

    """
    from ..core import Basic

    if isinstance(e, Basic):
        return e.count(Basic)
    if is_iterable(e):
        return 1 + sum(_nodes(ei) for ei in e)
    if isinstance(e, dict):
        return 1 + sum(_nodes(k) + _nodes(v) for k, v in e.items())
    return 1


def ordered(seq, keys=None, default=True, warn=False):
    """Return an iterator of the seq where keys are used to break ties in
    a conservative fashion: if, after applying a key, there are no ties
    then no other keys will be computed.

    Two default keys will be applied if 1) keys are not provided or 2) the
    given keys don't resolve all ties (but only if `default` is True). The
    two keys are `_nodes` (which places smaller expressions before large) and
    `default_sort_key` which (if the `sort_key` for an object is defined
    properly) should resolve any ties.

    If ``warn`` is True then an error will be raised if there were no
    keys remaining to break ties. This can be used if it was expected that
    there should be no ties between items that are not identical.

    Examples
    ========

    The count_ops is not sufficient to break ties in this list and the first
    two items appear in their original order (i.e. the sorting is stable):

    >>> list(ordered([y + 2, x + 2, x**2 + y + 3],
    ...              count_ops, default=False, warn=False))
    [y + 2, x + 2, x**2 + y + 3]

    The default_sort_key allows the tie to be broken:

    >>> list(ordered([y + 2, x + 2, x**2 + y + 3]))
    [x + 2, y + 2, x**2 + y + 3]

    Here, sequences are sorted by length, then sum:

    >>> seq, keys = [[[1, 2, 1], [0, 3, 1], [1, 1, 3], [2], [1]],
    ...              [lambda x: len(x), lambda x: sum(x)]]
    >>> list(ordered(seq, keys, default=False, warn=False))
    [[1], [2], [1, 2, 1], [0, 3, 1], [1, 1, 3]]

    If ``warn`` is True, an error will be raised if there were not
    enough keys to break ties:

    >>> list(ordered(seq, keys, default=False, warn=True))
    Traceback (most recent call last):
    ...
    ValueError: not enough keys to break ties

    Notes
    =====

    The decorated sort is one of the fastest ways to sort a sequence for
    which special item comparison is desired: the sequence is decorated,
    sorted on the basis of the decoration (e.g. making all letters lower
    case) and then undecorated. If one wants to break ties for items that
    have the same decorated value, a second key can be used. But if the
    second key is expensive to compute then it is inefficient to decorate
    all items with both keys: only those items having identical first key
    values need to be decorated. This function applies keys successively
    only when needed to break ties. By yielding an iterator, use of the
    tie-breaker is delayed as long as possible.

    This function is best used in cases when use of the first key is
    expected to be a good hashing function; if there are no unique hashes
    from application of a key then that key should not have been used. The
    exception, however, is that even if there are many collisions, if the
    first group is small and one does not need to process all items in the
    list then time will not be wasted sorting what one was not interested
    in. For example, if one were looking for the minimum in a list and
    there were several criteria used to define the sort order, then this
    function would be good at returning that quickly if the first group
    of candidates is small relative to the number of items being processed.

    """
    d = collections.defaultdict(list)
    if keys:
        if not isinstance(keys, (list, tuple)):
            keys = [keys]
        keys = list(keys)
        f = keys.pop(0)
        for a in seq:
            d[f(a)].append(a)
    else:
        if not default:
            raise ValueError('if default=False then keys must be provided')
        d[None].extend(seq)

    for k in sorted(d):
        if len(d[k]) > 1:
            if keys:
                d[k] = ordered(d[k], keys, default, warn)
            elif default:
                d[k] = ordered(d[k], (_nodes, default_sort_key,),
                               default=False, warn=warn)
            elif warn:
                u = list(uniq(d[k]))
                if len(u) > 1:
                    raise ValueError(f'not enough keys to break ties: {u}')
        yield from d[k]
        d.pop(k)


def default_sort_key(item, order=None):
    """
    Return a key that can be used for sorting.

    The key has the structure:

    (class_key, (len(args), args), exponent.sort_key(), coefficient)

    This key is supplied by the sort_key routine of Basic objects when
    ``item`` is a Basic object or an object (other than a string) that
    sympifies to a Basic object. Otherwise, this function produces the
    key.

    The ``order`` argument is passed along to the sort_key routine and is
    used to determine how the terms *within* an expression are ordered.
    (See examples below) ``order`` options are: 'lex', 'grlex', 'grevlex',
    and reversed values of the same (e.g. 'rev-lex'). The default order
    value is None (which translates to 'lex').

    Examples
    ========

    >>> from diofant.core.function import UndefinedFunction

    The following are equivalent ways of getting the key for an object:

    >>> x.sort_key() == default_sort_key(x)
    True

    Here are some examples of the key that is produced:

    >>> default_sort_key(UndefinedFunction('f'))
    ((0, 0, 'UndefinedFunction'), (1, ('f',)), ((1, 0, 'Number'),
        (0, ()), (), 1), 1)
    >>> default_sort_key('1')
    ((0, 0, 'str'), (1, ('1',)), ((1, 0, 'Number'), (0, ()), (), 1), 1)
    >>> default_sort_key(Integer(1))
    ((1, 0, 'Number'), (0, ()), (), 1)
    >>> default_sort_key(2)
    ((1, 0, 'Number'), (0, ()), (), 2)

    While sort_key is a method only defined for Diofant objects,
    default_sort_key will accept anything as an argument so it is
    more robust as a sorting key. For the following, using key=
    lambda i: i.sort_key() would fail because 2 doesn't have a sort_key
    method; that's why default_sort_key is used. Note, that it also
    handles sympification of non-string items likes ints:

    >>> a = [2, I, -I]
    >>> sorted(a, key=default_sort_key)
    [2, -I, I]

    The returned key can be used anywhere that a key can be specified for
    a function, e.g. sort, min, max, etc...:

    >>> a.sort(key=default_sort_key)
    >>> a[0]
    2
    >>> min(a, key=default_sort_key)
    2

    Notes
    =====

    The key returned is useful for getting items into a canonical order
    that will be the same across platforms. It is not directly useful for
    sorting lists of expressions:

    >>> a, b = x, 1/x

    Since ``a`` has only 1 term, its value of sort_key is unaffected by
    ``order``:

    >>> a.sort_key() == a.sort_key('rev-lex')
    True

    If ``a`` and ``b`` are combined then the key will differ because there
    are terms that can be ordered:

    >>> eq = a + b
    >>> eq.sort_key() == eq.sort_key('rev-lex')
    False
    >>> eq.as_ordered_terms()
    [x, 1/x]
    >>> eq.as_ordered_terms('rev-lex')
    [1/x, x]

    But since the keys for each of these terms are independent of ``order``'s
    value, they don't sort differently when they appear separately in a list:

    >>> sorted(eq.args, key=default_sort_key)
    [1/x, x]
    >>> sorted(eq.args, key=lambda i: default_sort_key(i, order='rev-lex'))
    [1/x, x]

    The order of terms obtained when using these keys is the order that would
    be obtained if those terms were *factors* in a product.

    Although it is useful for quickly putting expressions in canonical order,
    it does not sort expressions based on their complexity defined by the
    number of operations, power of variables and others:

    >>> sorted([sin(x)*cos(x), sin(x)], key=default_sort_key)
    [sin(x)*cos(x), sin(x)]
    >>> sorted([x, x**2, sqrt(x), x**3], key=default_sort_key)
    [sqrt(x), x, x**2, x**3]

    See Also
    ========

    ordered
    diofant.core.expr.Expr.as_ordered_factors
    diofant.core.expr.Expr.as_ordered_terms

    """
    from ..core import Basic, Integer, SympifyError, sympify

    if isinstance(item, Basic):
        return item.sort_key(order=order)

    if is_iterable(item, exclude=(str,)):
        if isinstance(item, dict):
            args = item.items()
            unordered = True
        elif isinstance(item, set):
            args = item
            unordered = True
        else:
            # e.g. tuple, list
            args = list(item)
            unordered = False

        args = [default_sort_key(arg, order=order) for arg in args]

        if unordered:  # e.g. dict, set
            args = sorted(args)

        cls_index, args = 10, (len(args), tuple(args))
    else:
        if not isinstance(item, str):
            try:
                item = sympify(item)
            except SympifyError:
                # e.g. lambda x: x
                pass
            else:
                if isinstance(item, Basic):
                    # e.g int -> Integer
                    return default_sort_key(item)
                # e.g. UndefinedFunction

        # e.g. str
        cls_index, args = 0, (1, (str(item),))

    return (cls_index, 0, item.__class__.__name__
            ), args, Integer(1).sort_key(), Integer(1)


class NotIterable:
    """
    Use this as mixin when creating a class which is not supposed to return
    true when is_iterable() is called on its instances. I.e. avoid infinite loop
    when calling e.g. list() on the instance

    """


def is_iterable(i, exclude=(str, dict, NotIterable)):
    """
    Return a boolean indicating whether ``i`` is Diofant iterable.
    True also indicates that the iterator is finite, i.e. you e.g.
    call list(...) on the instance.

    When Diofant is working with iterables, it is almost always assuming
    that the iterable is not a string or a mapping, so those are excluded
    by default. If you want a pure Python definition, make exclude=None. To
    exclude multiple items, pass them as a tuple.

    See Also
    ========

    is_sequence

    Examples
    ========

    >>> things = [[1], (1,), {1}, Tuple(1), (j for j in [1, 2]), {1: 2}, '1', 1]
    >>> for i in things:
    ...     print(f'{is_iterable(i)} {type(i)}')
    True <... 'list'>
    True <... 'tuple'>
    True <... 'set'>
    True <class 'diofant.core.containers.Tuple'>
    True <... 'generator'>
    False <... 'dict'>
    False <... 'str'>
    False <... 'int'>

    >>> is_iterable({}, exclude=None)
    True
    >>> is_iterable({}, exclude=str)
    True
    >>> is_iterable('no', exclude=str)
    False

    """
    try:
        iter(i)
    except TypeError:
        return False
    if exclude:
        return not isinstance(i, exclude)
    return True


def is_sequence(i, include=None):
    """
    Return a boolean indicating whether ``i`` is a sequence in the Diofant
    sense. If anything that fails the test below should be included as
    being a sequence for your application, set 'include' to that object's
    type; multiple types should be passed as a tuple of types.

    Note: although generators can generate a sequence, they often need special
    handling to make sure their elements are captured before the generator is
    exhausted, so these are not included by default in the definition of a
    sequence.

    See Also
    ========

    is_iterable

    Examples
    ========

    >>> from types import GeneratorType
    >>> is_sequence([])
    True
    >>> is_sequence(set())
    False
    >>> is_sequence('abc')
    False
    >>> is_sequence('abc', include=str)
    True
    >>> generator = (c for c in 'abc')
    >>> is_sequence(generator)
    False
    >>> is_sequence(generator, include=(str, GeneratorType))
    True

    """
    return (hasattr(i, '__getitem__') and
            is_iterable(i) or
            bool(include) and
            isinstance(i, include))