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

1""" 

2Algorithms and classes to support enumerative combinatorics. 

3 

4Currently just multiset partitions, but more could be added. 

5 

6Terminology (following Knuth, algorithm 7.1.2.5M TAOCP) 

7*multiset* aaabbcccc has a *partition* aaabc | bccc 

8 

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.) 

14 

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. 

18 

19The algorithm deserves some explanation. 

20 

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. 

35 

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. 

39 

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. 

49 

50Knuth's algorithm uses two main operations on parts: 

51 

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. 

61 

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. 

77 

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. 

82 

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. 

89 

90""" 

91 

92class PartComponent: 

93 """Internal class used in support of the multiset partitions 

94 enumerators and the associated visitor functions. 

95 

96 Represents one component of one part of the current partition. 

97 

98 A stack of these, plus an auxiliary frame array, f, represents a 

99 partition of the multiset. 

100 

101 Knuth's pseudocode makes c, u, and v separate arrays. 

102 """ 

103 

104 __slots__ = ('c', 'u', 'v') 

105 

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. 

115 

116 def __repr__(self): 

117 "for debug/algorithm animation purposes" 

118 return 'c:%d u:%d v:%d' % (self.c, self.u, self.v) 

119 

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) 

126 

127 def __ne__(self, other): 

128 """Defined for consistency with __eq__""" 

129 return not self == other 

130 

131 

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. 

137 

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. 

148 

149 Parameters 

150 ========== 

151 

152 multiplicities 

153 list of integer multiplicities of the components of the multiset. 

154 

155 Yields 

156 ====== 

157 

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. 

163 

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: 

167 

168 f 

169 is a frame array, which is used to divide pstack into parts. 

170 

171 lpart 

172 points to the base of the topmost part. 

173 

174 pstack 

175 is an array of PartComponent objects. 

176 

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. 

185 

186 Examples 

187 ======== 

188 

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']]] 

200 

201 See Also 

202 ======== 

203 

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. 

209 

210 """ 

211 

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) 

217 

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. 

221 

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) 

227 

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] 

235 

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 

242 

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 

264 

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 

274 

275 # M4 Visit a partition 

276 state = [f, lpart, pstack] 

277 yield state 

278 

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 

297 

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). 

302 

303 

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. 

310 

311 Examples 

312 ======== 

313 

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. 

317 

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 

339 

340 

341def list_visitor(state, components): 

342 """Return a list of lists to represent the partition. 

343 

344 Examples 

345 ======== 

346 

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 

358 

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) 

366 

367 return partition 

368 

369 

370class MultisetPartitionTraverser(): 

371 """ 

372 Has methods to ``enumerate`` and ``count`` the partitions of a multiset. 

373 

374 This implements a refactored and extended version of Knuth's algorithm 

375 7.1.2.5M [AOCP]_." 

376 

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``. 

380 

381 Examples 

382 ======== 

383 

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 

390 

391 See Also 

392 ======== 

393 

394 multiset_partitions_taocp 

395 sympy.utilities.iterables.multiset_partitions 

396 

397 References 

398 ========== 

399 

400 .. [AOCP] Algorithm 7.1.2.5M in Volume 4A, Combinatoral Algorithms, 

401 Part 1, of The Art of Computer Programming, by Donald Knuth. 

402 

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. 

408 

409 .. [Yorgey] Generating Multiset Partitions, Brent Yorgey, The 

410 Monad.Reader, Issue 8, September 2007. 

411 

412 """ 

413 

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 = [] 

428 

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** 

432 

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 

436 

437 if not hasattr(self, 'dp_map'): 

438 self.dp_map = {} 

439 

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)) 

452 

453 # 

454 # Helper methods for enumeration 

455 # 

456 def _initialize_enumeration(self, multiplicities): 

457 """Allocates and initializes the partition stack. 

458 

459 This is called from the enumeration/counting routines, so 

460 there is no need to call it separately.""" 

461 

462 num_components = len(multiplicities) 

463 # cardinality is the total number of elements, whether or not distinct 

464 cardinality = sum(multiplicities) 

465 

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) 

471 

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] 

478 

479 self.f[0] = 0 

480 self.f[1] = num_components 

481 self.lpart = 0 

482 

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. 

490 

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. 

495 

496 Parameters 

497 ========== 

498 

499 part 

500 The part, represented as a list of PartComponent objects, 

501 which is to be decremented. 

502 

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 

514 

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. 

520 

521 Parameters 

522 ========== 

523 

524 part 

525 part to be decremented (topmost part on the stack) 

526 

527 ub 

528 the maximum number of parts allowed in a partition 

529 returned by the calling traversal. 

530 

531 Notes 

532 ===== 

533 

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. 

544 

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: 

548 

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. 

552 

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. 

559 

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 

575 

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 

582 

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 

596 

597 def decrement_part_large(self, part, amt, lb): 

598 """Decrements part, while respecting size constraint. 

599 

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. 

605 

606 Returns True iff the part was successfully decremented. 

607 

608 Parameters 

609 ========== 

610 

611 part 

612 part to be decremented (topmost part on the stack) 

613 

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. 

619 

620 lb 

621 The partitions produced by the calling enumeration must 

622 have more parts than this value. 

623 

624 """ 

625 

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 

633 

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 

644 

645 deficit = min_unalloc - (total_mult - total_alloc) 

646 if deficit <= 0: 

647 return True 

648 

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 

663 

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. 

667 

668 Parameters 

669 ========== 

670 

671 part 

672 part to be decremented (topmost part on the stack) 

673 

674 ub 

675 the maximum number of parts allowed in a partition 

676 returned by the calling traversal. 

677 

678 lb 

679 The partitions produced by the calling enumeration must 

680 have more parts than this value. 

681 

682 Notes 

683 ===== 

684 

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 """ 

690 

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) 

699 

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. 

703 

704 Notes 

705 ===== 

706 

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. 

711 

712 This call does nothing (and returns False) if the current top 

713 part has no unallocated multiplicity. 

714 

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 

719 

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 

744 

745 def top_part(self): 

746 """Return current top part on the stack, as a slice of pstack. 

747 

748 """ 

749 return self.pstack[self.f[self.lpart]:self.f[self.lpart + 1]] 

750 

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. 

755 

756 Examples 

757 ======== 

758 

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']]] 

773 

774 See Also 

775 ======== 

776 

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. 

782 

783 """ 

784 self._initialize_enumeration(multiplicities) 

785 while True: 

786 while self.spread_part_multiplicity(): 

787 pass 

788 

789 # M4 Visit a partition 

790 state = [self.f, self.lpart, self.pstack] 

791 yield state 

792 

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 

799 

800 def enum_small(self, multiplicities, ub): 

801 """Enumerate multiset partitions with no more than ``ub`` parts. 

802 

803 Equivalent to enum_range(multiplicities, 0, ub) 

804 

805 Parameters 

806 ========== 

807 

808 multiplicities 

809 list of multiplicities of the components of the multiset. 

810 

811 ub 

812 Maximum number of parts 

813 

814 Examples 

815 ======== 

816 

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']]] 

827 

828 The implementation is based, in part, on the answer given to 

829 exercise 69, in Knuth [AOCP]_. 

830 

831 See Also 

832 ======== 

833 

834 enum_all, enum_large, enum_range 

835 

836 """ 

837 

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 

856 

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") 

866 

867 def enum_large(self, multiplicities, lb): 

868 """Enumerate the partitions of a multiset with lb < num(parts) 

869 

870 Equivalent to enum_range(multiplicities, lb, sum(multiplicities)) 

871 

872 Parameters 

873 ========== 

874 

875 multiplicities 

876 list of multiplicities of the components of the multiset. 

877 

878 lb 

879 Number of parts in the partition must be greater than 

880 this lower bound. 

881 

882 

883 Examples 

884 ======== 

885 

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']]] 

895 

896 See Also 

897 ======== 

898 

899 enum_all, enum_small, enum_range 

900 

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 

915 

916 # M4 Visit a partition 

917 if good_partition: 

918 state = [self.f, self.lpart, self.pstack] 

919 yield state 

920 

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 

927 

928 def enum_range(self, multiplicities, lb, ub): 

929 

930 """Enumerate the partitions of a multiset with 

931 ``lb < num(parts) <= ub``. 

932 

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. 

936 

937 Examples 

938 ======== 

939 

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']]] 

949 

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 

974 

975 # M4 Visit a partition 

976 if good_partition: 

977 state = [self.f, self.lpart, self.pstack] 

978 yield state 

979 

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") 

989 

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``. 

993 

994 Primarily for comparison purposes. It follows the same path as 

995 enumerate, and counts, rather than generates, the partitions. 

996 

997 See Also 

998 ======== 

999 

1000 count_partitions 

1001 Has the same calling interface, but is much faster. 

1002 

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 

1010 

1011 # M4 Visit (count) a partition 

1012 self.pcount += 1 

1013 

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 

1020 

1021 def count_partitions(self, multiplicities): 

1022 """Returns the number of partitions of a multiset whose components 

1023 have the multiplicities given in ``multiplicities``. 

1024 

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. 

1036 

1037 Examples 

1038 ======== 

1039 

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 

1047 

1048 Notes 

1049 ===== 

1050 

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. 

1062 

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. 

1068 

1069 This method follows a computational approach which is similar 

1070 to the hypothetical memoized recursive function, but with two 

1071 differences: 

1072 

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.) 

1079 

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. 

1084 

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 

1093 

1094 # dp_stack is list of lists of (part_key, start_count) pairs 

1095 self.dp_stack = [] 

1096 

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. 

1110 

1111 self.pcount += (self.dp_map[pkey] - 1) 

1112 self.lpart -= 1 

1113 break 

1114 else: 

1115 self.dp_stack.append([(pkey, self.pcount), ]) 

1116 

1117 # M4 count a leaf partition 

1118 self.pcount += 1 

1119 

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 

1128 

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),) 

1134 

1135 

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.) 

1141 

1142 Notes 

1143 ===== 

1144 

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. 

1149 

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)