partitions.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745
  1. from sympy.core import Basic, Dict, sympify, Tuple
  2. from sympy.core.numbers import Integer
  3. from sympy.core.sorting import default_sort_key
  4. from sympy.core.sympify import _sympify
  5. from sympy.functions.combinatorial.numbers import bell
  6. from sympy.matrices import zeros
  7. from sympy.sets.sets import FiniteSet, Union
  8. from sympy.utilities.iterables import flatten, group
  9. from sympy.utilities.misc import as_int
  10. from collections import defaultdict
  11. class Partition(FiniteSet):
  12. """
  13. This class represents an abstract partition.
  14. A partition is a set of disjoint sets whose union equals a given set.
  15. See Also
  16. ========
  17. sympy.utilities.iterables.partitions,
  18. sympy.utilities.iterables.multiset_partitions
  19. """
  20. _rank = None
  21. _partition = None
  22. def __new__(cls, *partition):
  23. """
  24. Generates a new partition object.
  25. This method also verifies if the arguments passed are
  26. valid and raises a ValueError if they are not.
  27. Examples
  28. ========
  29. Creating Partition from Python lists:
  30. >>> from sympy.combinatorics import Partition
  31. >>> a = Partition([1, 2], [3])
  32. >>> a
  33. Partition({3}, {1, 2})
  34. >>> a.partition
  35. [[1, 2], [3]]
  36. >>> len(a)
  37. 2
  38. >>> a.members
  39. (1, 2, 3)
  40. Creating Partition from Python sets:
  41. >>> Partition({1, 2, 3}, {4, 5})
  42. Partition({4, 5}, {1, 2, 3})
  43. Creating Partition from SymPy finite sets:
  44. >>> from sympy import FiniteSet
  45. >>> a = FiniteSet(1, 2, 3)
  46. >>> b = FiniteSet(4, 5)
  47. >>> Partition(a, b)
  48. Partition({4, 5}, {1, 2, 3})
  49. """
  50. args = []
  51. dups = False
  52. for arg in partition:
  53. if isinstance(arg, list):
  54. as_set = set(arg)
  55. if len(as_set) < len(arg):
  56. dups = True
  57. break # error below
  58. arg = as_set
  59. args.append(_sympify(arg))
  60. if not all(isinstance(part, FiniteSet) for part in args):
  61. raise ValueError(
  62. "Each argument to Partition should be " \
  63. "a list, set, or a FiniteSet")
  64. # sort so we have a canonical reference for RGS
  65. U = Union(*args)
  66. if dups or len(U) < sum(len(arg) for arg in args):
  67. raise ValueError("Partition contained duplicate elements.")
  68. obj = FiniteSet.__new__(cls, *args)
  69. obj.members = tuple(U)
  70. obj.size = len(U)
  71. return obj
  72. def sort_key(self, order=None):
  73. """Return a canonical key that can be used for sorting.
  74. Ordering is based on the size and sorted elements of the partition
  75. and ties are broken with the rank.
  76. Examples
  77. ========
  78. >>> from sympy import default_sort_key
  79. >>> from sympy.combinatorics import Partition
  80. >>> from sympy.abc import x
  81. >>> a = Partition([1, 2])
  82. >>> b = Partition([3, 4])
  83. >>> c = Partition([1, x])
  84. >>> d = Partition(list(range(4)))
  85. >>> l = [d, b, a + 1, a, c]
  86. >>> l.sort(key=default_sort_key); l
  87. [Partition({1, 2}), Partition({1}, {2}), Partition({1, x}), Partition({3, 4}), Partition({0, 1, 2, 3})]
  88. """
  89. if order is None:
  90. members = self.members
  91. else:
  92. members = tuple(sorted(self.members,
  93. key=lambda w: default_sort_key(w, order)))
  94. return tuple(map(default_sort_key, (self.size, members, self.rank)))
  95. @property
  96. def partition(self):
  97. """Return partition as a sorted list of lists.
  98. Examples
  99. ========
  100. >>> from sympy.combinatorics import Partition
  101. >>> Partition([1], [2, 3]).partition
  102. [[1], [2, 3]]
  103. """
  104. if self._partition is None:
  105. self._partition = sorted([sorted(p, key=default_sort_key)
  106. for p in self.args])
  107. return self._partition
  108. def __add__(self, other):
  109. """
  110. Return permutation whose rank is ``other`` greater than current rank,
  111. (mod the maximum rank for the set).
  112. Examples
  113. ========
  114. >>> from sympy.combinatorics import Partition
  115. >>> a = Partition([1, 2], [3])
  116. >>> a.rank
  117. 1
  118. >>> (a + 1).rank
  119. 2
  120. >>> (a + 100).rank
  121. 1
  122. """
  123. other = as_int(other)
  124. offset = self.rank + other
  125. result = RGS_unrank((offset) %
  126. RGS_enum(self.size),
  127. self.size)
  128. return Partition.from_rgs(result, self.members)
  129. def __sub__(self, other):
  130. """
  131. Return permutation whose rank is ``other`` less than current rank,
  132. (mod the maximum rank for the set).
  133. Examples
  134. ========
  135. >>> from sympy.combinatorics import Partition
  136. >>> a = Partition([1, 2], [3])
  137. >>> a.rank
  138. 1
  139. >>> (a - 1).rank
  140. 0
  141. >>> (a - 100).rank
  142. 1
  143. """
  144. return self.__add__(-other)
  145. def __le__(self, other):
  146. """
  147. Checks if a partition is less than or equal to
  148. the other based on rank.
  149. Examples
  150. ========
  151. >>> from sympy.combinatorics import Partition
  152. >>> a = Partition([1, 2], [3, 4, 5])
  153. >>> b = Partition([1], [2, 3], [4], [5])
  154. >>> a.rank, b.rank
  155. (9, 34)
  156. >>> a <= a
  157. True
  158. >>> a <= b
  159. True
  160. """
  161. return self.sort_key() <= sympify(other).sort_key()
  162. def __lt__(self, other):
  163. """
  164. Checks if a partition is less than the other.
  165. Examples
  166. ========
  167. >>> from sympy.combinatorics import Partition
  168. >>> a = Partition([1, 2], [3, 4, 5])
  169. >>> b = Partition([1], [2, 3], [4], [5])
  170. >>> a.rank, b.rank
  171. (9, 34)
  172. >>> a < b
  173. True
  174. """
  175. return self.sort_key() < sympify(other).sort_key()
  176. @property
  177. def rank(self):
  178. """
  179. Gets the rank of a partition.
  180. Examples
  181. ========
  182. >>> from sympy.combinatorics import Partition
  183. >>> a = Partition([1, 2], [3], [4, 5])
  184. >>> a.rank
  185. 13
  186. """
  187. if self._rank is not None:
  188. return self._rank
  189. self._rank = RGS_rank(self.RGS)
  190. return self._rank
  191. @property
  192. def RGS(self):
  193. """
  194. Returns the "restricted growth string" of the partition.
  195. Explanation
  196. ===========
  197. The RGS is returned as a list of indices, L, where L[i] indicates
  198. the block in which element i appears. For example, in a partition
  199. of 3 elements (a, b, c) into 2 blocks ([c], [a, b]) the RGS is
  200. [1, 1, 0]: "a" is in block 1, "b" is in block 1 and "c" is in block 0.
  201. Examples
  202. ========
  203. >>> from sympy.combinatorics import Partition
  204. >>> a = Partition([1, 2], [3], [4, 5])
  205. >>> a.members
  206. (1, 2, 3, 4, 5)
  207. >>> a.RGS
  208. (0, 0, 1, 2, 2)
  209. >>> a + 1
  210. Partition({3}, {4}, {5}, {1, 2})
  211. >>> _.RGS
  212. (0, 0, 1, 2, 3)
  213. """
  214. rgs = {}
  215. partition = self.partition
  216. for i, part in enumerate(partition):
  217. for j in part:
  218. rgs[j] = i
  219. return tuple([rgs[i] for i in sorted(
  220. [i for p in partition for i in p], key=default_sort_key)])
  221. @classmethod
  222. def from_rgs(self, rgs, elements):
  223. """
  224. Creates a set partition from a restricted growth string.
  225. Explanation
  226. ===========
  227. The indices given in rgs are assumed to be the index
  228. of the element as given in elements *as provided* (the
  229. elements are not sorted by this routine). Block numbering
  230. starts from 0. If any block was not referenced in ``rgs``
  231. an error will be raised.
  232. Examples
  233. ========
  234. >>> from sympy.combinatorics import Partition
  235. >>> Partition.from_rgs([0, 1, 2, 0, 1], list('abcde'))
  236. Partition({c}, {a, d}, {b, e})
  237. >>> Partition.from_rgs([0, 1, 2, 0, 1], list('cbead'))
  238. Partition({e}, {a, c}, {b, d})
  239. >>> a = Partition([1, 4], [2], [3, 5])
  240. >>> Partition.from_rgs(a.RGS, a.members)
  241. Partition({2}, {1, 4}, {3, 5})
  242. """
  243. if len(rgs) != len(elements):
  244. raise ValueError('mismatch in rgs and element lengths')
  245. max_elem = max(rgs) + 1
  246. partition = [[] for i in range(max_elem)]
  247. j = 0
  248. for i in rgs:
  249. partition[i].append(elements[j])
  250. j += 1
  251. if not all(p for p in partition):
  252. raise ValueError('some blocks of the partition were empty.')
  253. return Partition(*partition)
  254. class IntegerPartition(Basic):
  255. """
  256. This class represents an integer partition.
  257. Explanation
  258. ===========
  259. In number theory and combinatorics, a partition of a positive integer,
  260. ``n``, also called an integer partition, is a way of writing ``n`` as a
  261. list of positive integers that sum to n. Two partitions that differ only
  262. in the order of summands are considered to be the same partition; if order
  263. matters then the partitions are referred to as compositions. For example,
  264. 4 has five partitions: [4], [3, 1], [2, 2], [2, 1, 1], and [1, 1, 1, 1];
  265. the compositions [1, 2, 1] and [1, 1, 2] are the same as partition
  266. [2, 1, 1].
  267. See Also
  268. ========
  269. sympy.utilities.iterables.partitions,
  270. sympy.utilities.iterables.multiset_partitions
  271. References
  272. ==========
  273. .. [1] https://en.wikipedia.org/wiki/Partition_%28number_theory%29
  274. """
  275. _dict = None
  276. _keys = None
  277. def __new__(cls, partition, integer=None):
  278. """
  279. Generates a new IntegerPartition object from a list or dictionary.
  280. Explantion
  281. ==========
  282. The partition can be given as a list of positive integers or a
  283. dictionary of (integer, multiplicity) items. If the partition is
  284. preceded by an integer an error will be raised if the partition
  285. does not sum to that given integer.
  286. Examples
  287. ========
  288. >>> from sympy.combinatorics.partitions import IntegerPartition
  289. >>> a = IntegerPartition([5, 4, 3, 1, 1])
  290. >>> a
  291. IntegerPartition(14, (5, 4, 3, 1, 1))
  292. >>> print(a)
  293. [5, 4, 3, 1, 1]
  294. >>> IntegerPartition({1:3, 2:1})
  295. IntegerPartition(5, (2, 1, 1, 1))
  296. If the value that the partition should sum to is given first, a check
  297. will be made to see n error will be raised if there is a discrepancy:
  298. >>> IntegerPartition(10, [5, 4, 3, 1])
  299. Traceback (most recent call last):
  300. ...
  301. ValueError: The partition is not valid
  302. """
  303. if integer is not None:
  304. integer, partition = partition, integer
  305. if isinstance(partition, (dict, Dict)):
  306. _ = []
  307. for k, v in sorted(list(partition.items()), reverse=True):
  308. if not v:
  309. continue
  310. k, v = as_int(k), as_int(v)
  311. _.extend([k]*v)
  312. partition = tuple(_)
  313. else:
  314. partition = tuple(sorted(map(as_int, partition), reverse=True))
  315. sum_ok = False
  316. if integer is None:
  317. integer = sum(partition)
  318. sum_ok = True
  319. else:
  320. integer = as_int(integer)
  321. if not sum_ok and sum(partition) != integer:
  322. raise ValueError("Partition did not add to %s" % integer)
  323. if any(i < 1 for i in partition):
  324. raise ValueError("All integer summands must be greater than one")
  325. obj = Basic.__new__(cls, Integer(integer), Tuple(*partition))
  326. obj.partition = list(partition)
  327. obj.integer = integer
  328. return obj
  329. def prev_lex(self):
  330. """Return the previous partition of the integer, n, in lexical order,
  331. wrapping around to [1, ..., 1] if the partition is [n].
  332. Examples
  333. ========
  334. >>> from sympy.combinatorics.partitions import IntegerPartition
  335. >>> p = IntegerPartition([4])
  336. >>> print(p.prev_lex())
  337. [3, 1]
  338. >>> p.partition > p.prev_lex().partition
  339. True
  340. """
  341. d = defaultdict(int)
  342. d.update(self.as_dict())
  343. keys = self._keys
  344. if keys == [1]:
  345. return IntegerPartition({self.integer: 1})
  346. if keys[-1] != 1:
  347. d[keys[-1]] -= 1
  348. if keys[-1] == 2:
  349. d[1] = 2
  350. else:
  351. d[keys[-1] - 1] = d[1] = 1
  352. else:
  353. d[keys[-2]] -= 1
  354. left = d[1] + keys[-2]
  355. new = keys[-2]
  356. d[1] = 0
  357. while left:
  358. new -= 1
  359. if left - new >= 0:
  360. d[new] += left//new
  361. left -= d[new]*new
  362. return IntegerPartition(self.integer, d)
  363. def next_lex(self):
  364. """Return the next partition of the integer, n, in lexical order,
  365. wrapping around to [n] if the partition is [1, ..., 1].
  366. Examples
  367. ========
  368. >>> from sympy.combinatorics.partitions import IntegerPartition
  369. >>> p = IntegerPartition([3, 1])
  370. >>> print(p.next_lex())
  371. [4]
  372. >>> p.partition < p.next_lex().partition
  373. True
  374. """
  375. d = defaultdict(int)
  376. d.update(self.as_dict())
  377. key = self._keys
  378. a = key[-1]
  379. if a == self.integer:
  380. d.clear()
  381. d[1] = self.integer
  382. elif a == 1:
  383. if d[a] > 1:
  384. d[a + 1] += 1
  385. d[a] -= 2
  386. else:
  387. b = key[-2]
  388. d[b + 1] += 1
  389. d[1] = (d[b] - 1)*b
  390. d[b] = 0
  391. else:
  392. if d[a] > 1:
  393. if len(key) == 1:
  394. d.clear()
  395. d[a + 1] = 1
  396. d[1] = self.integer - a - 1
  397. else:
  398. a1 = a + 1
  399. d[a1] += 1
  400. d[1] = d[a]*a - a1
  401. d[a] = 0
  402. else:
  403. b = key[-2]
  404. b1 = b + 1
  405. d[b1] += 1
  406. need = d[b]*b + d[a]*a - b1
  407. d[a] = d[b] = 0
  408. d[1] = need
  409. return IntegerPartition(self.integer, d)
  410. def as_dict(self):
  411. """Return the partition as a dictionary whose keys are the
  412. partition integers and the values are the multiplicity of that
  413. integer.
  414. Examples
  415. ========
  416. >>> from sympy.combinatorics.partitions import IntegerPartition
  417. >>> IntegerPartition([1]*3 + [2] + [3]*4).as_dict()
  418. {1: 3, 2: 1, 3: 4}
  419. """
  420. if self._dict is None:
  421. groups = group(self.partition, multiple=False)
  422. self._keys = [g[0] for g in groups]
  423. self._dict = dict(groups)
  424. return self._dict
  425. @property
  426. def conjugate(self):
  427. """
  428. Computes the conjugate partition of itself.
  429. Examples
  430. ========
  431. >>> from sympy.combinatorics.partitions import IntegerPartition
  432. >>> a = IntegerPartition([6, 3, 3, 2, 1])
  433. >>> a.conjugate
  434. [5, 4, 3, 1, 1, 1]
  435. """
  436. j = 1
  437. temp_arr = list(self.partition) + [0]
  438. k = temp_arr[0]
  439. b = [0]*k
  440. while k > 0:
  441. while k > temp_arr[j]:
  442. b[k - 1] = j
  443. k -= 1
  444. j += 1
  445. return b
  446. def __lt__(self, other):
  447. """Return True if self is less than other when the partition
  448. is listed from smallest to biggest.
  449. Examples
  450. ========
  451. >>> from sympy.combinatorics.partitions import IntegerPartition
  452. >>> a = IntegerPartition([3, 1])
  453. >>> a < a
  454. False
  455. >>> b = a.next_lex()
  456. >>> a < b
  457. True
  458. >>> a == b
  459. False
  460. """
  461. return list(reversed(self.partition)) < list(reversed(other.partition))
  462. def __le__(self, other):
  463. """Return True if self is less than other when the partition
  464. is listed from smallest to biggest.
  465. Examples
  466. ========
  467. >>> from sympy.combinatorics.partitions import IntegerPartition
  468. >>> a = IntegerPartition([4])
  469. >>> a <= a
  470. True
  471. """
  472. return list(reversed(self.partition)) <= list(reversed(other.partition))
  473. def as_ferrers(self, char='#'):
  474. """
  475. Prints the ferrer diagram of a partition.
  476. Examples
  477. ========
  478. >>> from sympy.combinatorics.partitions import IntegerPartition
  479. >>> print(IntegerPartition([1, 1, 5]).as_ferrers())
  480. #####
  481. #
  482. #
  483. """
  484. return "\n".join([char*i for i in self.partition])
  485. def __str__(self):
  486. return str(list(self.partition))
  487. def random_integer_partition(n, seed=None):
  488. """
  489. Generates a random integer partition summing to ``n`` as a list
  490. of reverse-sorted integers.
  491. Examples
  492. ========
  493. >>> from sympy.combinatorics.partitions import random_integer_partition
  494. For the following, a seed is given so a known value can be shown; in
  495. practice, the seed would not be given.
  496. >>> random_integer_partition(100, seed=[1, 1, 12, 1, 2, 1, 85, 1])
  497. [85, 12, 2, 1]
  498. >>> random_integer_partition(10, seed=[1, 2, 3, 1, 5, 1])
  499. [5, 3, 1, 1]
  500. >>> random_integer_partition(1)
  501. [1]
  502. """
  503. from sympy.core.random import _randint
  504. n = as_int(n)
  505. if n < 1:
  506. raise ValueError('n must be a positive integer')
  507. randint = _randint(seed)
  508. partition = []
  509. while (n > 0):
  510. k = randint(1, n)
  511. mult = randint(1, n//k)
  512. partition.append((k, mult))
  513. n -= k*mult
  514. partition.sort(reverse=True)
  515. partition = flatten([[k]*m for k, m in partition])
  516. return partition
  517. def RGS_generalized(m):
  518. """
  519. Computes the m + 1 generalized unrestricted growth strings
  520. and returns them as rows in matrix.
  521. Examples
  522. ========
  523. >>> from sympy.combinatorics.partitions import RGS_generalized
  524. >>> RGS_generalized(6)
  525. Matrix([
  526. [ 1, 1, 1, 1, 1, 1, 1],
  527. [ 1, 2, 3, 4, 5, 6, 0],
  528. [ 2, 5, 10, 17, 26, 0, 0],
  529. [ 5, 15, 37, 77, 0, 0, 0],
  530. [ 15, 52, 151, 0, 0, 0, 0],
  531. [ 52, 203, 0, 0, 0, 0, 0],
  532. [203, 0, 0, 0, 0, 0, 0]])
  533. """
  534. d = zeros(m + 1)
  535. for i in range(0, m + 1):
  536. d[0, i] = 1
  537. for i in range(1, m + 1):
  538. for j in range(m):
  539. if j <= m - i:
  540. d[i, j] = j * d[i - 1, j] + d[i - 1, j + 1]
  541. else:
  542. d[i, j] = 0
  543. return d
  544. def RGS_enum(m):
  545. """
  546. RGS_enum computes the total number of restricted growth strings
  547. possible for a superset of size m.
  548. Examples
  549. ========
  550. >>> from sympy.combinatorics.partitions import RGS_enum
  551. >>> from sympy.combinatorics import Partition
  552. >>> RGS_enum(4)
  553. 15
  554. >>> RGS_enum(5)
  555. 52
  556. >>> RGS_enum(6)
  557. 203
  558. We can check that the enumeration is correct by actually generating
  559. the partitions. Here, the 15 partitions of 4 items are generated:
  560. >>> a = Partition(list(range(4)))
  561. >>> s = set()
  562. >>> for i in range(20):
  563. ... s.add(a)
  564. ... a += 1
  565. ...
  566. >>> assert len(s) == 15
  567. """
  568. if (m < 1):
  569. return 0
  570. elif (m == 1):
  571. return 1
  572. else:
  573. return bell(m)
  574. def RGS_unrank(rank, m):
  575. """
  576. Gives the unranked restricted growth string for a given
  577. superset size.
  578. Examples
  579. ========
  580. >>> from sympy.combinatorics.partitions import RGS_unrank
  581. >>> RGS_unrank(14, 4)
  582. [0, 1, 2, 3]
  583. >>> RGS_unrank(0, 4)
  584. [0, 0, 0, 0]
  585. """
  586. if m < 1:
  587. raise ValueError("The superset size must be >= 1")
  588. if rank < 0 or RGS_enum(m) <= rank:
  589. raise ValueError("Invalid arguments")
  590. L = [1] * (m + 1)
  591. j = 1
  592. D = RGS_generalized(m)
  593. for i in range(2, m + 1):
  594. v = D[m - i, j]
  595. cr = j*v
  596. if cr <= rank:
  597. L[i] = j + 1
  598. rank -= cr
  599. j += 1
  600. else:
  601. L[i] = int(rank / v + 1)
  602. rank %= v
  603. return [x - 1 for x in L[1:]]
  604. def RGS_rank(rgs):
  605. """
  606. Computes the rank of a restricted growth string.
  607. Examples
  608. ========
  609. >>> from sympy.combinatorics.partitions import RGS_rank, RGS_unrank
  610. >>> RGS_rank([0, 1, 2, 1, 3])
  611. 42
  612. >>> RGS_rank(RGS_unrank(4, 7))
  613. 4
  614. """
  615. rgs_size = len(rgs)
  616. rank = 0
  617. D = RGS_generalized(rgs_size)
  618. for i in range(1, rgs_size):
  619. n = len(rgs[(i + 1):])
  620. m = max(rgs[0:i])
  621. rank += D[n, m + 1] * rgs[i]
  622. return rank