diff --git a/sympy/utilities/iterables.py b/sympy/utilities/iterables.py index 144713d4fe47..a0eb4cece751 100644 --- a/sympy/utilities/iterables.py +++ b/sympy/utilities/iterables.py @@ -3,7 +3,6 @@ chain, combinations, combinations_with_replacement, cycle, islice, permutations, product ) - # For backwards compatibility from itertools import product as cartes # noqa: F401 from operator import gt @@ -2706,6 +2705,130 @@ def runs(seq, op=gt): return cycles +def sequence_partitions(l, n, /): + r"""Returns the partition of sequence $l$ into $n$ bins + + Explanation + =========== + + Given the sequence $l_1 \cdots l_m \in V^+$ where + $V^+$ is the Kleene plus of $V$ + + The set of $n$ partitions of $l$ is defined as: + + .. math:: + \{(s_1, \cdots, s_n) | s_1 \in V^+, \cdots, s_n \in V^+, + s_1 \cdots s_n = l_1 \cdots l_m\} + + Parameters + ========== + + l : Sequence[T] + A nonempty sequence of any Python objects + + n : int + A positive integer + + Yields + ====== + + out : list[Sequence[T]] + A list of sequences with concatenation equals $l$. + This should conform with the type of $l$. + + Examples + ======== + + >>> from sympy.utilities.iterables import sequence_partitions + >>> for out in sequence_partitions([1, 2, 3, 4], 2): + ... print(out) + [[1], [2, 3, 4]] + [[1, 2], [3, 4]] + [[1, 2, 3], [4]] + + Notes + ===== + + This is modified version of EnricoGiampieri's partition generator + from https://stackoverflow.com/questions/13131491/ + + See Also + ======== + + sequence_partitions_empty + """ + # Asserting l is nonempty is done only for sanity check + if n == 1 and l: + yield [l] + return + for i in range(1, len(l)): + for part in sequence_partitions(l[i:], n - 1): + yield [l[:i]] + part + + +def sequence_partitions_empty(l, n, /): + r"""Returns the partition of sequence $l$ into $n$ bins with + empty sequence + + Explanation + =========== + + Given the sequence $l_1 \cdots l_m \in V^*$ where + $V^*$ is the Kleene star of $V$ + + The set of $n$ partitions of $l$ is defined as: + + .. math:: + \{(s_1, \cdots, s_n) | s_1 \in V^*, \cdots, s_n \in V^*, + s_1 \cdots s_n = l_1 \cdots l_m\} + + There are more combinations than :func:`sequence_partitions` because + empty sequence can fill everywhere, so we try to provide different + utility for this. + + Parameters + ========== + + l : Sequence[T] + A sequence of any Python objects (can be possibly empty) + + n : int + A positive integer + + Yields + ====== + + out : list[Sequence[T]] + A list of sequences with concatenation equals $l$. + This should conform with the type of $l$. + + Examples + ======== + + >>> from sympy.utilities.iterables import sequence_partitions_empty + >>> for out in sequence_partitions_empty([1, 2, 3, 4], 2): + ... print(out) + [[], [1, 2, 3, 4]] + [[1], [2, 3, 4]] + [[1, 2], [3, 4]] + [[1, 2, 3], [4]] + [[1, 2, 3, 4], []] + + See Also + ======== + + sequence_partitions + """ + if n < 1: + return + if n == 1: + yield [l] + return + for i in range(0, len(l) + 1): + for part in sequence_partitions_empty(l[i:], n - 1): + yield [l[:i]] + part + + def kbins(l, k, ordered=None): """ Return sequence ``l`` partitioned into ``k`` bins. @@ -2787,24 +2910,12 @@ def kbins(l, k, ordered=None): partitions, multiset_partitions """ - def partition(lista, bins): - # EnricoGiampieri's partition generator from - # https://stackoverflow.com/questions/13131491/ - # partition-n-items-into-k-bins-in-python-lazily - if len(lista) == 1 or bins == 1: - yield [lista] - elif len(lista) > 1 and bins > 1: - for i in range(1, len(lista)): - for part in partition(lista[i:], bins - 1): - if len([lista[:i]] + part) == bins: - yield [lista[:i]] + part - if ordered is None: - yield from partition(l, k) + yield from sequence_partitions(l, k) elif ordered == 11: for pl in multiset_permutations(l): pl = list(pl) - yield from partition(pl, k) + yield from sequence_partitions(pl, k) elif ordered == 00: yield from multiset_partitions(l, k) elif ordered == 10: diff --git a/sympy/utilities/tests/test_iterables.py b/sympy/utilities/tests/test_iterables.py index 77c6e7b1ee80..8e527a447e75 100644 --- a/sympy/utilities/tests/test_iterables.py +++ b/sympy/utilities/tests/test_iterables.py @@ -19,7 +19,8 @@ prefixes, reshape, rotate_left, rotate_right, runs, sift, strongly_connected_components, subsets, take, topological_sort, unflatten, uniq, variations, ordered_partitions, rotations, is_palindromic, iterable, - NotIterable, multiset_derangements) + NotIterable, multiset_derangements, + sequence_partitions, sequence_partitions_empty) from sympy.utilities.enumerative import ( factoring_visitor, multiset_partitions_taocp ) @@ -885,3 +886,51 @@ class Test6(Test5): _iterable = False assert iterable(Test6()) is False + + +def test_sequence_partitions(): + assert list(sequence_partitions([1], 1)) == [[[1]]] + assert list(sequence_partitions([1, 2], 1)) == [[[1, 2]]] + assert list(sequence_partitions([1, 2], 2)) == [[[1], [2]]] + assert list(sequence_partitions([1, 2, 3], 1)) == [[[1, 2, 3]]] + assert list(sequence_partitions([1, 2, 3], 2)) == \ + [[[1], [2, 3]], [[1, 2], [3]]] + assert list(sequence_partitions([1, 2, 3], 3)) == [[[1], [2], [3]]] + + # Exceptional cases + assert list(sequence_partitions([], 0)) == [] + assert list(sequence_partitions([], 1)) == [] + assert list(sequence_partitions([1, 2], 0)) == [] + assert list(sequence_partitions([1, 2], 3)) == [] + + +def test_sequence_partitions_empty(): + assert list(sequence_partitions_empty([], 1)) == [[[]]] + assert list(sequence_partitions_empty([], 2)) == [[[], []]] + assert list(sequence_partitions_empty([], 3)) == [[[], [], []]] + assert list(sequence_partitions_empty([1], 1)) == [[[1]]] + assert list(sequence_partitions_empty([1], 2)) == [[[], [1]], [[1], []]] + assert list(sequence_partitions_empty([1], 3)) == \ + [[[], [], [1]], [[], [1], []], [[1], [], []]] + assert list(sequence_partitions_empty([1, 2], 1)) == [[[1, 2]]] + assert list(sequence_partitions_empty([1, 2], 2)) == \ + [[[], [1, 2]], [[1], [2]], [[1, 2], []]] + assert list(sequence_partitions_empty([1, 2], 3)) == [ + [[], [], [1, 2]], [[], [1], [2]], [[], [1, 2], []], + [[1], [], [2]], [[1], [2], []], [[1, 2], [], []] + ] + assert list(sequence_partitions_empty([1, 2, 3], 1)) == [[[1, 2, 3]]] + assert list(sequence_partitions_empty([1, 2, 3], 2)) == \ + [[[], [1, 2, 3]], [[1], [2, 3]], [[1, 2], [3]], [[1, 2, 3], []]] + assert list(sequence_partitions_empty([1, 2, 3], 3)) == [ + [[], [], [1, 2, 3]], [[], [1], [2, 3]], + [[], [1, 2], [3]], [[], [1, 2, 3], []], + [[1], [], [2, 3]], [[1], [2], [3]], + [[1], [2, 3], []], [[1, 2], [], [3]], + [[1, 2], [3], []], [[1, 2, 3], [], []] + ] + + # Exceptional cases + assert list(sequence_partitions([], 0)) == [] + assert list(sequence_partitions([1], 0)) == [] + assert list(sequence_partitions([1, 2], 0)) == []