Coverage for /usr/lib/python3/dist-packages/sympy/utilities/enumerative.py: 8%
319 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1"""
2Algorithms and classes to support enumerative combinatorics.
4Currently just multiset partitions, but more could be added.
6Terminology (following Knuth, algorithm 7.1.2.5M TAOCP)
7*multiset* aaabbcccc has a *partition* aaabc | bccc
9The submultisets, aaabc and bccc of the partition are called
10*parts*, or sometimes *vectors*. (Knuth notes that multiset
11partitions can be thought of as partitions of vectors of integers,
12where the ith element of the vector gives the multiplicity of
13element i.)
15The values a, b and c are *components* of the multiset. These
16correspond to elements of a set, but in a multiset can be present
17with a multiplicity greater than 1.
19The algorithm deserves some explanation.
21Think of the part aaabc from the multiset above. If we impose an
22ordering on the components of the multiset, we can represent a part
23with a vector, in which the value of the first element of the vector
24corresponds to the multiplicity of the first component in that
25part. Thus, aaabc can be represented by the vector [3, 1, 1]. We
26can also define an ordering on parts, based on the lexicographic
27ordering of the vector (leftmost vector element, i.e., the element
28with the smallest component number, is the most significant), so
29that [3, 1, 1] > [3, 1, 0] and [3, 1, 1] > [2, 1, 4]. The ordering
30on parts can be extended to an ordering on partitions: First, sort
31the parts in each partition, left-to-right in decreasing order. Then
32partition A is greater than partition B if A's leftmost/greatest
33part is greater than B's leftmost part. If the leftmost parts are
34equal, compare the second parts, and so on.
36In this ordering, the greatest partition of a given multiset has only
37one part. The least partition is the one in which the components
38are spread out, one per part.
40The enumeration algorithms in this file yield the partitions of the
41argument multiset in decreasing order. The main data structure is a
42stack of parts, corresponding to the current partition. An
43important invariant is that the parts on the stack are themselves in
44decreasing order. This data structure is decremented to find the
45next smaller partition. Most often, decrementing the partition will
46only involve adjustments to the smallest parts at the top of the
47stack, much as adjacent integers *usually* differ only in their last
48few digits.
50Knuth's algorithm uses two main operations on parts:
52Decrement - change the part so that it is smaller in the
53 (vector) lexicographic order, but reduced by the smallest amount possible.
54 For example, if the multiset has vector [5,
55 3, 1], and the bottom/greatest part is [4, 2, 1], this part would
56 decrement to [4, 2, 0], while [4, 0, 0] would decrement to [3, 3,
57 1]. A singleton part is never decremented -- [1, 0, 0] is not
58 decremented to [0, 3, 1]. Instead, the decrement operator needs
59 to fail for this case. In Knuth's pseudocode, the decrement
60 operator is step m5.
62Spread unallocated multiplicity - Once a part has been decremented,
63 it cannot be the rightmost part in the partition. There is some
64 multiplicity that has not been allocated, and new parts must be
65 created above it in the stack to use up this multiplicity. To
66 maintain the invariant that the parts on the stack are in
67 decreasing order, these new parts must be less than or equal to
68 the decremented part.
69 For example, if the multiset is [5, 3, 1], and its most
70 significant part has just been decremented to [5, 3, 0], the
71 spread operation will add a new part so that the stack becomes
72 [[5, 3, 0], [0, 0, 1]]. If the most significant part (for the
73 same multiset) has been decremented to [2, 0, 0] the stack becomes
74 [[2, 0, 0], [2, 0, 0], [1, 3, 1]]. In the pseudocode, the spread
75 operation for one part is step m2. The complete spread operation
76 is a loop of steps m2 and m3.
78In order to facilitate the spread operation, Knuth stores, for each
79component of each part, not just the multiplicity of that component
80in the part, but also the total multiplicity available for this
81component in this part or any lesser part above it on the stack.
83One added twist is that Knuth does not represent the part vectors as
84arrays. Instead, he uses a sparse representation, in which a
85component of a part is represented as a component number (c), plus
86the multiplicity of the component in that part (v) as well as the
87total multiplicity available for that component (u). This saves
88time that would be spent skipping over zeros.
90"""
92class PartComponent:
93 """Internal class used in support of the multiset partitions
94 enumerators and the associated visitor functions.
96 Represents one component of one part of the current partition.
98 A stack of these, plus an auxiliary frame array, f, represents a
99 partition of the multiset.
101 Knuth's pseudocode makes c, u, and v separate arrays.
102 """
104 __slots__ = ('c', 'u', 'v')
106 def __init__(self):
107 self.c = 0 # Component number
108 self.u = 0 # The as yet unpartitioned amount in component c
109 # *before* it is allocated by this triple
110 self.v = 0 # Amount of c component in the current part
111 # (v<=u). An invariant of the representation is
112 # that the next higher triple for this component
113 # (if there is one) will have a value of u-v in
114 # its u attribute.
116 def __repr__(self):
117 "for debug/algorithm animation purposes"
118 return 'c:%d u:%d v:%d' % (self.c, self.u, self.v)
120 def __eq__(self, other):
121 """Define value oriented equality, which is useful for testers"""
122 return (isinstance(other, self.__class__) and
123 self.c == other.c and
124 self.u == other.u and
125 self.v == other.v)
127 def __ne__(self, other):
128 """Defined for consistency with __eq__"""
129 return not self == other
132# This function tries to be a faithful implementation of algorithm
133# 7.1.2.5M in Volume 4A, Combinatoral Algorithms, Part 1, of The Art
134# of Computer Programming, by Donald Knuth. This includes using
135# (mostly) the same variable names, etc. This makes for rather
136# low-level Python.
138# Changes from Knuth's pseudocode include
139# - use PartComponent struct/object instead of 3 arrays
140# - make the function a generator
141# - map (with some difficulty) the GOTOs to Python control structures.
142# - Knuth uses 1-based numbering for components, this code is 0-based
143# - renamed variable l to lpart.
144# - flag variable x takes on values True/False instead of 1/0
145#
146def multiset_partitions_taocp(multiplicities):
147 """Enumerates partitions of a multiset.
149 Parameters
150 ==========
152 multiplicities
153 list of integer multiplicities of the components of the multiset.
155 Yields
156 ======
158 state
159 Internal data structure which encodes a particular partition.
160 This output is then usually processed by a visitor function
161 which combines the information from this data structure with
162 the components themselves to produce an actual partition.
164 Unless they wish to create their own visitor function, users will
165 have little need to look inside this data structure. But, for
166 reference, it is a 3-element list with components:
168 f
169 is a frame array, which is used to divide pstack into parts.
171 lpart
172 points to the base of the topmost part.
174 pstack
175 is an array of PartComponent objects.
177 The ``state`` output offers a peek into the internal data
178 structures of the enumeration function. The client should
179 treat this as read-only; any modification of the data
180 structure will cause unpredictable (and almost certainly
181 incorrect) results. Also, the components of ``state`` are
182 modified in place at each iteration. Hence, the visitor must
183 be called at each loop iteration. Accumulating the ``state``
184 instances and processing them later will not work.
186 Examples
187 ========
189 >>> from sympy.utilities.enumerative import list_visitor
190 >>> from sympy.utilities.enumerative import multiset_partitions_taocp
191 >>> # variables components and multiplicities represent the multiset 'abb'
192 >>> components = 'ab'
193 >>> multiplicities = [1, 2]
194 >>> states = multiset_partitions_taocp(multiplicities)
195 >>> list(list_visitor(state, components) for state in states)
196 [[['a', 'b', 'b']],
197 [['a', 'b'], ['b']],
198 [['a'], ['b', 'b']],
199 [['a'], ['b'], ['b']]]
201 See Also
202 ========
204 sympy.utilities.iterables.multiset_partitions: Takes a multiset
205 as input and directly yields multiset partitions. It
206 dispatches to a number of functions, including this one, for
207 implementation. Most users will find it more convenient to
208 use than multiset_partitions_taocp.
210 """
212 # Important variables.
213 # m is the number of components, i.e., number of distinct elements
214 m = len(multiplicities)
215 # n is the cardinality, total number of elements whether or not distinct
216 n = sum(multiplicities)
218 # The main data structure, f segments pstack into parts. See
219 # list_visitor() for example code indicating how this internal
220 # state corresponds to a partition.
222 # Note: allocation of space for stack is conservative. Knuth's
223 # exercise 7.2.1.5.68 gives some indication of how to tighten this
224 # bound, but this is not implemented.
225 pstack = [PartComponent() for i in range(n * m + 1)]
226 f = [0] * (n + 1)
228 # Step M1 in Knuth (Initialize)
229 # Initial state - entire multiset in one part.
230 for j in range(m):
231 ps = pstack[j]
232 ps.c = j
233 ps.u = multiplicities[j]
234 ps.v = multiplicities[j]
236 # Other variables
237 f[0] = 0
238 a = 0
239 lpart = 0
240 f[1] = m
241 b = m # in general, current stack frame is from a to b - 1
243 while True:
244 while True:
245 # Step M2 (Subtract v from u)
246 j = a
247 k = b
248 x = False
249 while j < b:
250 pstack[k].u = pstack[j].u - pstack[j].v
251 if pstack[k].u == 0:
252 x = True
253 elif not x:
254 pstack[k].c = pstack[j].c
255 pstack[k].v = min(pstack[j].v, pstack[k].u)
256 x = pstack[k].u < pstack[j].v
257 k = k + 1
258 else: # x is True
259 pstack[k].c = pstack[j].c
260 pstack[k].v = pstack[k].u
261 k = k + 1
262 j = j + 1
263 # Note: x is True iff v has changed
265 # Step M3 (Push if nonzero.)
266 if k > b:
267 a = b
268 b = k
269 lpart = lpart + 1
270 f[lpart + 1] = b
271 # Return to M2
272 else:
273 break # Continue to M4
275 # M4 Visit a partition
276 state = [f, lpart, pstack]
277 yield state
279 # M5 (Decrease v)
280 while True:
281 j = b-1
282 while (pstack[j].v == 0):
283 j = j - 1
284 if j == a and pstack[j].v == 1:
285 # M6 (Backtrack)
286 if lpart == 0:
287 return
288 lpart = lpart - 1
289 b = a
290 a = f[lpart]
291 # Return to M5
292 else:
293 pstack[j].v = pstack[j].v - 1
294 for k in range(j + 1, b):
295 pstack[k].v = pstack[k].u
296 break # GOTO M2
298# --------------- Visitor functions for multiset partitions ---------------
299# A visitor takes the partition state generated by
300# multiset_partitions_taocp or other enumerator, and produces useful
301# output (such as the actual partition).
304def factoring_visitor(state, primes):
305 """Use with multiset_partitions_taocp to enumerate the ways a
306 number can be expressed as a product of factors. For this usage,
307 the exponents of the prime factors of a number are arguments to
308 the partition enumerator, while the corresponding prime factors
309 are input here.
311 Examples
312 ========
314 To enumerate the factorings of a number we can think of the elements of the
315 partition as being the prime factors and the multiplicities as being their
316 exponents.
318 >>> from sympy.utilities.enumerative import factoring_visitor
319 >>> from sympy.utilities.enumerative import multiset_partitions_taocp
320 >>> from sympy import factorint
321 >>> primes, multiplicities = zip(*factorint(24).items())
322 >>> primes
323 (2, 3)
324 >>> multiplicities
325 (3, 1)
326 >>> states = multiset_partitions_taocp(multiplicities)
327 >>> list(factoring_visitor(state, primes) for state in states)
328 [[24], [8, 3], [12, 2], [4, 6], [4, 2, 3], [6, 2, 2], [2, 2, 2, 3]]
329 """
330 f, lpart, pstack = state
331 factoring = []
332 for i in range(lpart + 1):
333 factor = 1
334 for ps in pstack[f[i]: f[i + 1]]:
335 if ps.v > 0:
336 factor *= primes[ps.c] ** ps.v
337 factoring.append(factor)
338 return factoring
341def list_visitor(state, components):
342 """Return a list of lists to represent the partition.
344 Examples
345 ========
347 >>> from sympy.utilities.enumerative import list_visitor
348 >>> from sympy.utilities.enumerative import multiset_partitions_taocp
349 >>> states = multiset_partitions_taocp([1, 2, 1])
350 >>> s = next(states)
351 >>> list_visitor(s, 'abc') # for multiset 'a b b c'
352 [['a', 'b', 'b', 'c']]
353 >>> s = next(states)
354 >>> list_visitor(s, [1, 2, 3]) # for multiset '1 2 2 3
355 [[1, 2, 2], [3]]
356 """
357 f, lpart, pstack = state
359 partition = []
360 for i in range(lpart+1):
361 part = []
362 for ps in pstack[f[i]:f[i+1]]:
363 if ps.v > 0:
364 part.extend([components[ps.c]] * ps.v)
365 partition.append(part)
367 return partition
370class MultisetPartitionTraverser():
371 """
372 Has methods to ``enumerate`` and ``count`` the partitions of a multiset.
374 This implements a refactored and extended version of Knuth's algorithm
375 7.1.2.5M [AOCP]_."
377 The enumeration methods of this class are generators and return
378 data structures which can be interpreted by the same visitor
379 functions used for the output of ``multiset_partitions_taocp``.
381 Examples
382 ========
384 >>> from sympy.utilities.enumerative import MultisetPartitionTraverser
385 >>> m = MultisetPartitionTraverser()
386 >>> m.count_partitions([4,4,4,2])
387 127750
388 >>> m.count_partitions([3,3,3])
389 686
391 See Also
392 ========
394 multiset_partitions_taocp
395 sympy.utilities.iterables.multiset_partitions
397 References
398 ==========
400 .. [AOCP] Algorithm 7.1.2.5M in Volume 4A, Combinatoral Algorithms,
401 Part 1, of The Art of Computer Programming, by Donald Knuth.
403 .. [Factorisatio] On a Problem of Oppenheim concerning
404 "Factorisatio Numerorum" E. R. Canfield, Paul Erdos, Carl
405 Pomerance, JOURNAL OF NUMBER THEORY, Vol. 17, No. 1. August
406 1983. See section 7 for a description of an algorithm
407 similar to Knuth's.
409 .. [Yorgey] Generating Multiset Partitions, Brent Yorgey, The
410 Monad.Reader, Issue 8, September 2007.
412 """
414 def __init__(self):
415 self.debug = False
416 # TRACING variables. These are useful for gathering
417 # statistics on the algorithm itself, but have no particular
418 # benefit to a user of the code.
419 self.k1 = 0
420 self.k2 = 0
421 self.p1 = 0
422 self.pstack = None
423 self.f = None
424 self.lpart = 0
425 self.discarded = 0
426 # dp_stack is list of lists of (part_key, start_count) pairs
427 self.dp_stack = []
429 # dp_map is map part_key-> count, where count represents the
430 # number of multiset which are descendants of a part with this
431 # key, **or any of its decrements**
433 # Thus, when we find a part in the map, we add its count
434 # value to the running total, cut off the enumeration, and
435 # backtrack
437 if not hasattr(self, 'dp_map'):
438 self.dp_map = {}
440 def db_trace(self, msg):
441 """Useful for understanding/debugging the algorithms. Not
442 generally activated in end-user code."""
443 if self.debug:
444 # XXX: animation_visitor is undefined... Clearly this does not
445 # work and was not tested. Previous code in comments below.
446 raise RuntimeError
447 #letters = 'abcdefghijklmnopqrstuvwxyz'
448 #state = [self.f, self.lpart, self.pstack]
449 #print("DBG:", msg,
450 # ["".join(part) for part in list_visitor(state, letters)],
451 # animation_visitor(state))
453 #
454 # Helper methods for enumeration
455 #
456 def _initialize_enumeration(self, multiplicities):
457 """Allocates and initializes the partition stack.
459 This is called from the enumeration/counting routines, so
460 there is no need to call it separately."""
462 num_components = len(multiplicities)
463 # cardinality is the total number of elements, whether or not distinct
464 cardinality = sum(multiplicities)
466 # pstack is the partition stack, which is segmented by
467 # f into parts.
468 self.pstack = [PartComponent() for i in
469 range(num_components * cardinality + 1)]
470 self.f = [0] * (cardinality + 1)
472 # Initial state - entire multiset in one part.
473 for j in range(num_components):
474 ps = self.pstack[j]
475 ps.c = j
476 ps.u = multiplicities[j]
477 ps.v = multiplicities[j]
479 self.f[0] = 0
480 self.f[1] = num_components
481 self.lpart = 0
483 # The decrement_part() method corresponds to step M5 in Knuth's
484 # algorithm. This is the base version for enum_all(). Modified
485 # versions of this method are needed if we want to restrict
486 # sizes of the partitions produced.
487 def decrement_part(self, part):
488 """Decrements part (a subrange of pstack), if possible, returning
489 True iff the part was successfully decremented.
491 If you think of the v values in the part as a multi-digit
492 integer (least significant digit on the right) this is
493 basically decrementing that integer, but with the extra
494 constraint that the leftmost digit cannot be decremented to 0.
496 Parameters
497 ==========
499 part
500 The part, represented as a list of PartComponent objects,
501 which is to be decremented.
503 """
504 plen = len(part)
505 for j in range(plen - 1, -1, -1):
506 if j == 0 and part[j].v > 1 or j > 0 and part[j].v > 0:
507 # found val to decrement
508 part[j].v -= 1
509 # Reset trailing parts back to maximum
510 for k in range(j + 1, plen):
511 part[k].v = part[k].u
512 return True
513 return False
515 # Version to allow number of parts to be bounded from above.
516 # Corresponds to (a modified) step M5.
517 def decrement_part_small(self, part, ub):
518 """Decrements part (a subrange of pstack), if possible, returning
519 True iff the part was successfully decremented.
521 Parameters
522 ==========
524 part
525 part to be decremented (topmost part on the stack)
527 ub
528 the maximum number of parts allowed in a partition
529 returned by the calling traversal.
531 Notes
532 =====
534 The goal of this modification of the ordinary decrement method
535 is to fail (meaning that the subtree rooted at this part is to
536 be skipped) when it can be proved that this part can only have
537 child partitions which are larger than allowed by ``ub``. If a
538 decision is made to fail, it must be accurate, otherwise the
539 enumeration will miss some partitions. But, it is OK not to
540 capture all the possible failures -- if a part is passed that
541 should not be, the resulting too-large partitions are filtered
542 by the enumeration one level up. However, as is usual in
543 constrained enumerations, failing early is advantageous.
545 The tests used by this method catch the most common cases,
546 although this implementation is by no means the last word on
547 this problem. The tests include:
549 1) ``lpart`` must be less than ``ub`` by at least 2. This is because
550 once a part has been decremented, the partition
551 will gain at least one child in the spread step.
553 2) If the leading component of the part is about to be
554 decremented, check for how many parts will be added in
555 order to use up the unallocated multiplicity in that
556 leading component, and fail if this number is greater than
557 allowed by ``ub``. (See code for the exact expression.) This
558 test is given in the answer to Knuth's problem 7.2.1.5.69.
560 3) If there is *exactly* enough room to expand the leading
561 component by the above test, check the next component (if
562 it exists) once decrementing has finished. If this has
563 ``v == 0``, this next component will push the expansion over the
564 limit by 1, so fail.
565 """
566 if self.lpart >= ub - 1:
567 self.p1 += 1 # increment to keep track of usefulness of tests
568 return False
569 plen = len(part)
570 for j in range(plen - 1, -1, -1):
571 # Knuth's mod, (answer to problem 7.2.1.5.69)
572 if j == 0 and (part[0].v - 1)*(ub - self.lpart) < part[0].u:
573 self.k1 += 1
574 return False
576 if j == 0 and part[j].v > 1 or j > 0 and part[j].v > 0:
577 # found val to decrement
578 part[j].v -= 1
579 # Reset trailing parts back to maximum
580 for k in range(j + 1, plen):
581 part[k].v = part[k].u
583 # Have now decremented part, but are we doomed to
584 # failure when it is expanded? Check one oddball case
585 # that turns out to be surprisingly common - exactly
586 # enough room to expand the leading component, but no
587 # room for the second component, which has v=0.
588 if (plen > 1 and part[1].v == 0 and
589 (part[0].u - part[0].v) ==
590 ((ub - self.lpart - 1) * part[0].v)):
591 self.k2 += 1
592 self.db_trace("Decrement fails test 3")
593 return False
594 return True
595 return False
597 def decrement_part_large(self, part, amt, lb):
598 """Decrements part, while respecting size constraint.
600 A part can have no children which are of sufficient size (as
601 indicated by ``lb``) unless that part has sufficient
602 unallocated multiplicity. When enforcing the size constraint,
603 this method will decrement the part (if necessary) by an
604 amount needed to ensure sufficient unallocated multiplicity.
606 Returns True iff the part was successfully decremented.
608 Parameters
609 ==========
611 part
612 part to be decremented (topmost part on the stack)
614 amt
615 Can only take values 0 or 1. A value of 1 means that the
616 part must be decremented, and then the size constraint is
617 enforced. A value of 0 means just to enforce the ``lb``
618 size constraint.
620 lb
621 The partitions produced by the calling enumeration must
622 have more parts than this value.
624 """
626 if amt == 1:
627 # In this case we always need to increment, *before*
628 # enforcing the "sufficient unallocated multiplicity"
629 # constraint. Easiest for this is just to call the
630 # regular decrement method.
631 if not self.decrement_part(part):
632 return False
634 # Next, perform any needed additional decrementing to respect
635 # "sufficient unallocated multiplicity" (or fail if this is
636 # not possible).
637 min_unalloc = lb - self.lpart
638 if min_unalloc <= 0:
639 return True
640 total_mult = sum(pc.u for pc in part)
641 total_alloc = sum(pc.v for pc in part)
642 if total_mult <= min_unalloc:
643 return False
645 deficit = min_unalloc - (total_mult - total_alloc)
646 if deficit <= 0:
647 return True
649 for i in range(len(part) - 1, -1, -1):
650 if i == 0:
651 if part[0].v > deficit:
652 part[0].v -= deficit
653 return True
654 else:
655 return False # This shouldn't happen, due to above check
656 else:
657 if part[i].v >= deficit:
658 part[i].v -= deficit
659 return True
660 else:
661 deficit -= part[i].v
662 part[i].v = 0
664 def decrement_part_range(self, part, lb, ub):
665 """Decrements part (a subrange of pstack), if possible, returning
666 True iff the part was successfully decremented.
668 Parameters
669 ==========
671 part
672 part to be decremented (topmost part on the stack)
674 ub
675 the maximum number of parts allowed in a partition
676 returned by the calling traversal.
678 lb
679 The partitions produced by the calling enumeration must
680 have more parts than this value.
682 Notes
683 =====
685 Combines the constraints of _small and _large decrement
686 methods. If returns success, part has been decremented at
687 least once, but perhaps by quite a bit more if needed to meet
688 the lb constraint.
689 """
691 # Constraint in the range case is just enforcing both the
692 # constraints from _small and _large cases. Note the 0 as the
693 # second argument to the _large call -- this is the signal to
694 # decrement only as needed to for constraint enforcement. The
695 # short circuiting and left-to-right order of the 'and'
696 # operator is important for this to work correctly.
697 return self.decrement_part_small(part, ub) and \
698 self.decrement_part_large(part, 0, lb)
700 def spread_part_multiplicity(self):
701 """Returns True if a new part has been created, and
702 adjusts pstack, f and lpart as needed.
704 Notes
705 =====
707 Spreads unallocated multiplicity from the current top part
708 into a new part created above the current on the stack. This
709 new part is constrained to be less than or equal to the old in
710 terms of the part ordering.
712 This call does nothing (and returns False) if the current top
713 part has no unallocated multiplicity.
715 """
716 j = self.f[self.lpart] # base of current top part
717 k = self.f[self.lpart + 1] # ub of current; potential base of next
718 base = k # save for later comparison
720 changed = False # Set to true when the new part (so far) is
721 # strictly less than (as opposed to less than
722 # or equal) to the old.
723 for j in range(self.f[self.lpart], self.f[self.lpart + 1]):
724 self.pstack[k].u = self.pstack[j].u - self.pstack[j].v
725 if self.pstack[k].u == 0:
726 changed = True
727 else:
728 self.pstack[k].c = self.pstack[j].c
729 if changed: # Put all available multiplicity in this part
730 self.pstack[k].v = self.pstack[k].u
731 else: # Still maintaining ordering constraint
732 if self.pstack[k].u < self.pstack[j].v:
733 self.pstack[k].v = self.pstack[k].u
734 changed = True
735 else:
736 self.pstack[k].v = self.pstack[j].v
737 k = k + 1
738 if k > base:
739 # Adjust for the new part on stack
740 self.lpart = self.lpart + 1
741 self.f[self.lpart + 1] = k
742 return True
743 return False
745 def top_part(self):
746 """Return current top part on the stack, as a slice of pstack.
748 """
749 return self.pstack[self.f[self.lpart]:self.f[self.lpart + 1]]
751 # Same interface and functionality as multiset_partitions_taocp(),
752 # but some might find this refactored version easier to follow.
753 def enum_all(self, multiplicities):
754 """Enumerate the partitions of a multiset.
756 Examples
757 ========
759 >>> from sympy.utilities.enumerative import list_visitor
760 >>> from sympy.utilities.enumerative import MultisetPartitionTraverser
761 >>> m = MultisetPartitionTraverser()
762 >>> states = m.enum_all([2,2])
763 >>> list(list_visitor(state, 'ab') for state in states)
764 [[['a', 'a', 'b', 'b']],
765 [['a', 'a', 'b'], ['b']],
766 [['a', 'a'], ['b', 'b']],
767 [['a', 'a'], ['b'], ['b']],
768 [['a', 'b', 'b'], ['a']],
769 [['a', 'b'], ['a', 'b']],
770 [['a', 'b'], ['a'], ['b']],
771 [['a'], ['a'], ['b', 'b']],
772 [['a'], ['a'], ['b'], ['b']]]
774 See Also
775 ========
777 multiset_partitions_taocp:
778 which provides the same result as this method, but is
779 about twice as fast. Hence, enum_all is primarily useful
780 for testing. Also see the function for a discussion of
781 states and visitors.
783 """
784 self._initialize_enumeration(multiplicities)
785 while True:
786 while self.spread_part_multiplicity():
787 pass
789 # M4 Visit a partition
790 state = [self.f, self.lpart, self.pstack]
791 yield state
793 # M5 (Decrease v)
794 while not self.decrement_part(self.top_part()):
795 # M6 (Backtrack)
796 if self.lpart == 0:
797 return
798 self.lpart -= 1
800 def enum_small(self, multiplicities, ub):
801 """Enumerate multiset partitions with no more than ``ub`` parts.
803 Equivalent to enum_range(multiplicities, 0, ub)
805 Parameters
806 ==========
808 multiplicities
809 list of multiplicities of the components of the multiset.
811 ub
812 Maximum number of parts
814 Examples
815 ========
817 >>> from sympy.utilities.enumerative import list_visitor
818 >>> from sympy.utilities.enumerative import MultisetPartitionTraverser
819 >>> m = MultisetPartitionTraverser()
820 >>> states = m.enum_small([2,2], 2)
821 >>> list(list_visitor(state, 'ab') for state in states)
822 [[['a', 'a', 'b', 'b']],
823 [['a', 'a', 'b'], ['b']],
824 [['a', 'a'], ['b', 'b']],
825 [['a', 'b', 'b'], ['a']],
826 [['a', 'b'], ['a', 'b']]]
828 The implementation is based, in part, on the answer given to
829 exercise 69, in Knuth [AOCP]_.
831 See Also
832 ========
834 enum_all, enum_large, enum_range
836 """
838 # Keep track of iterations which do not yield a partition.
839 # Clearly, we would like to keep this number small.
840 self.discarded = 0
841 if ub <= 0:
842 return
843 self._initialize_enumeration(multiplicities)
844 while True:
845 while self.spread_part_multiplicity():
846 self.db_trace('spread 1')
847 if self.lpart >= ub:
848 self.discarded += 1
849 self.db_trace(' Discarding')
850 self.lpart = ub - 2
851 break
852 else:
853 # M4 Visit a partition
854 state = [self.f, self.lpart, self.pstack]
855 yield state
857 # M5 (Decrease v)
858 while not self.decrement_part_small(self.top_part(), ub):
859 self.db_trace("Failed decrement, going to backtrack")
860 # M6 (Backtrack)
861 if self.lpart == 0:
862 return
863 self.lpart -= 1
864 self.db_trace("Backtracked to")
865 self.db_trace("decrement ok, about to expand")
867 def enum_large(self, multiplicities, lb):
868 """Enumerate the partitions of a multiset with lb < num(parts)
870 Equivalent to enum_range(multiplicities, lb, sum(multiplicities))
872 Parameters
873 ==========
875 multiplicities
876 list of multiplicities of the components of the multiset.
878 lb
879 Number of parts in the partition must be greater than
880 this lower bound.
883 Examples
884 ========
886 >>> from sympy.utilities.enumerative import list_visitor
887 >>> from sympy.utilities.enumerative import MultisetPartitionTraverser
888 >>> m = MultisetPartitionTraverser()
889 >>> states = m.enum_large([2,2], 2)
890 >>> list(list_visitor(state, 'ab') for state in states)
891 [[['a', 'a'], ['b'], ['b']],
892 [['a', 'b'], ['a'], ['b']],
893 [['a'], ['a'], ['b', 'b']],
894 [['a'], ['a'], ['b'], ['b']]]
896 See Also
897 ========
899 enum_all, enum_small, enum_range
901 """
902 self.discarded = 0
903 if lb >= sum(multiplicities):
904 return
905 self._initialize_enumeration(multiplicities)
906 self.decrement_part_large(self.top_part(), 0, lb)
907 while True:
908 good_partition = True
909 while self.spread_part_multiplicity():
910 if not self.decrement_part_large(self.top_part(), 0, lb):
911 # Failure here should be rare/impossible
912 self.discarded += 1
913 good_partition = False
914 break
916 # M4 Visit a partition
917 if good_partition:
918 state = [self.f, self.lpart, self.pstack]
919 yield state
921 # M5 (Decrease v)
922 while not self.decrement_part_large(self.top_part(), 1, lb):
923 # M6 (Backtrack)
924 if self.lpart == 0:
925 return
926 self.lpart -= 1
928 def enum_range(self, multiplicities, lb, ub):
930 """Enumerate the partitions of a multiset with
931 ``lb < num(parts) <= ub``.
933 In particular, if partitions with exactly ``k`` parts are
934 desired, call with ``(multiplicities, k - 1, k)``. This
935 method generalizes enum_all, enum_small, and enum_large.
937 Examples
938 ========
940 >>> from sympy.utilities.enumerative import list_visitor
941 >>> from sympy.utilities.enumerative import MultisetPartitionTraverser
942 >>> m = MultisetPartitionTraverser()
943 >>> states = m.enum_range([2,2], 1, 2)
944 >>> list(list_visitor(state, 'ab') for state in states)
945 [[['a', 'a', 'b'], ['b']],
946 [['a', 'a'], ['b', 'b']],
947 [['a', 'b', 'b'], ['a']],
948 [['a', 'b'], ['a', 'b']]]
950 """
951 # combine the constraints of the _large and _small
952 # enumerations.
953 self.discarded = 0
954 if ub <= 0 or lb >= sum(multiplicities):
955 return
956 self._initialize_enumeration(multiplicities)
957 self.decrement_part_large(self.top_part(), 0, lb)
958 while True:
959 good_partition = True
960 while self.spread_part_multiplicity():
961 self.db_trace("spread 1")
962 if not self.decrement_part_large(self.top_part(), 0, lb):
963 # Failure here - possible in range case?
964 self.db_trace(" Discarding (large cons)")
965 self.discarded += 1
966 good_partition = False
967 break
968 elif self.lpart >= ub:
969 self.discarded += 1
970 good_partition = False
971 self.db_trace(" Discarding small cons")
972 self.lpart = ub - 2
973 break
975 # M4 Visit a partition
976 if good_partition:
977 state = [self.f, self.lpart, self.pstack]
978 yield state
980 # M5 (Decrease v)
981 while not self.decrement_part_range(self.top_part(), lb, ub):
982 self.db_trace("Failed decrement, going to backtrack")
983 # M6 (Backtrack)
984 if self.lpart == 0:
985 return
986 self.lpart -= 1
987 self.db_trace("Backtracked to")
988 self.db_trace("decrement ok, about to expand")
990 def count_partitions_slow(self, multiplicities):
991 """Returns the number of partitions of a multiset whose elements
992 have the multiplicities given in ``multiplicities``.
994 Primarily for comparison purposes. It follows the same path as
995 enumerate, and counts, rather than generates, the partitions.
997 See Also
998 ========
1000 count_partitions
1001 Has the same calling interface, but is much faster.
1003 """
1004 # number of partitions so far in the enumeration
1005 self.pcount = 0
1006 self._initialize_enumeration(multiplicities)
1007 while True:
1008 while self.spread_part_multiplicity():
1009 pass
1011 # M4 Visit (count) a partition
1012 self.pcount += 1
1014 # M5 (Decrease v)
1015 while not self.decrement_part(self.top_part()):
1016 # M6 (Backtrack)
1017 if self.lpart == 0:
1018 return self.pcount
1019 self.lpart -= 1
1021 def count_partitions(self, multiplicities):
1022 """Returns the number of partitions of a multiset whose components
1023 have the multiplicities given in ``multiplicities``.
1025 For larger counts, this method is much faster than calling one
1026 of the enumerators and counting the result. Uses dynamic
1027 programming to cut down on the number of nodes actually
1028 explored. The dictionary used in order to accelerate the
1029 counting process is stored in the ``MultisetPartitionTraverser``
1030 object and persists across calls. If the user does not
1031 expect to call ``count_partitions`` for any additional
1032 multisets, the object should be cleared to save memory. On
1033 the other hand, the cache built up from one count run can
1034 significantly speed up subsequent calls to ``count_partitions``,
1035 so it may be advantageous not to clear the object.
1037 Examples
1038 ========
1040 >>> from sympy.utilities.enumerative import MultisetPartitionTraverser
1041 >>> m = MultisetPartitionTraverser()
1042 >>> m.count_partitions([9,8,2])
1043 288716
1044 >>> m.count_partitions([2,2])
1045 9
1046 >>> del m
1048 Notes
1049 =====
1051 If one looks at the workings of Knuth's algorithm M [AOCP]_, it
1052 can be viewed as a traversal of a binary tree of parts. A
1053 part has (up to) two children, the left child resulting from
1054 the spread operation, and the right child from the decrement
1055 operation. The ordinary enumeration of multiset partitions is
1056 an in-order traversal of this tree, and with the partitions
1057 corresponding to paths from the root to the leaves. The
1058 mapping from paths to partitions is a little complicated,
1059 since the partition would contain only those parts which are
1060 leaves or the parents of a spread link, not those which are
1061 parents of a decrement link.
1063 For counting purposes, it is sufficient to count leaves, and
1064 this can be done with a recursive in-order traversal. The
1065 number of leaves of a subtree rooted at a particular part is a
1066 function only of that part itself, so memoizing has the
1067 potential to speed up the counting dramatically.
1069 This method follows a computational approach which is similar
1070 to the hypothetical memoized recursive function, but with two
1071 differences:
1073 1) This method is iterative, borrowing its structure from the
1074 other enumerations and maintaining an explicit stack of
1075 parts which are in the process of being counted. (There
1076 may be multisets which can be counted reasonably quickly by
1077 this implementation, but which would overflow the default
1078 Python recursion limit with a recursive implementation.)
1080 2) Instead of using the part data structure directly, a more
1081 compact key is constructed. This saves space, but more
1082 importantly coalesces some parts which would remain
1083 separate with physical keys.
1085 Unlike the enumeration functions, there is currently no _range
1086 version of count_partitions. If someone wants to stretch
1087 their brain, it should be possible to construct one by
1088 memoizing with a histogram of counts rather than a single
1089 count, and combining the histograms.
1090 """
1091 # number of partitions so far in the enumeration
1092 self.pcount = 0
1094 # dp_stack is list of lists of (part_key, start_count) pairs
1095 self.dp_stack = []
1097 self._initialize_enumeration(multiplicities)
1098 pkey = part_key(self.top_part())
1099 self.dp_stack.append([(pkey, 0), ])
1100 while True:
1101 while self.spread_part_multiplicity():
1102 pkey = part_key(self.top_part())
1103 if pkey in self.dp_map:
1104 # Already have a cached value for the count of the
1105 # subtree rooted at this part. Add it to the
1106 # running counter, and break out of the spread
1107 # loop. The -1 below is to compensate for the
1108 # leaf that this code path would otherwise find,
1109 # and which gets incremented for below.
1111 self.pcount += (self.dp_map[pkey] - 1)
1112 self.lpart -= 1
1113 break
1114 else:
1115 self.dp_stack.append([(pkey, self.pcount), ])
1117 # M4 count a leaf partition
1118 self.pcount += 1
1120 # M5 (Decrease v)
1121 while not self.decrement_part(self.top_part()):
1122 # M6 (Backtrack)
1123 for key, oldcount in self.dp_stack.pop():
1124 self.dp_map[key] = self.pcount - oldcount
1125 if self.lpart == 0:
1126 return self.pcount
1127 self.lpart -= 1
1129 # At this point have successfully decremented the part on
1130 # the stack and it does not appear in the cache. It needs
1131 # to be added to the list at the top of dp_stack
1132 pkey = part_key(self.top_part())
1133 self.dp_stack[-1].append((pkey, self.pcount),)
1136def part_key(part):
1137 """Helper for MultisetPartitionTraverser.count_partitions that
1138 creates a key for ``part``, that only includes information which can
1139 affect the count for that part. (Any irrelevant information just
1140 reduces the effectiveness of dynamic programming.)
1142 Notes
1143 =====
1145 This member function is a candidate for future exploration. There
1146 are likely symmetries that can be exploited to coalesce some
1147 ``part_key`` values, and thereby save space and improve
1148 performance.
1150 """
1151 # The component number is irrelevant for counting partitions, so
1152 # leave it out of the memo key.
1153 rval = []
1154 for ps in part:
1155 rval.append(ps.u)
1156 rval.append(ps.v)
1157 return tuple(rval)