subsets.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619
  1. from itertools import combinations
  2. from sympy.combinatorics.graycode import GrayCode
  3. class Subset():
  4. """
  5. Represents a basic subset object.
  6. Explanation
  7. ===========
  8. We generate subsets using essentially two techniques,
  9. binary enumeration and lexicographic enumeration.
  10. The Subset class takes two arguments, the first one
  11. describes the initial subset to consider and the second
  12. describes the superset.
  13. Examples
  14. ========
  15. >>> from sympy.combinatorics import Subset
  16. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  17. >>> a.next_binary().subset
  18. ['b']
  19. >>> a.prev_binary().subset
  20. ['c']
  21. """
  22. _rank_binary = None
  23. _rank_lex = None
  24. _rank_graycode = None
  25. _subset = None
  26. _superset = None
  27. def __new__(cls, subset, superset):
  28. """
  29. Default constructor.
  30. It takes the ``subset`` and its ``superset`` as its parameters.
  31. Examples
  32. ========
  33. >>> from sympy.combinatorics import Subset
  34. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  35. >>> a.subset
  36. ['c', 'd']
  37. >>> a.superset
  38. ['a', 'b', 'c', 'd']
  39. >>> a.size
  40. 2
  41. """
  42. if len(subset) > len(superset):
  43. raise ValueError('Invalid arguments have been provided. The '
  44. 'superset must be larger than the subset.')
  45. for elem in subset:
  46. if elem not in superset:
  47. raise ValueError('The superset provided is invalid as it does '
  48. 'not contain the element {}'.format(elem))
  49. obj = object.__new__(cls)
  50. obj._subset = subset
  51. obj._superset = superset
  52. return obj
  53. def __eq__(self, other):
  54. """Return a boolean indicating whether a == b on the basis of
  55. whether both objects are of the class Subset and if the values
  56. of the subset and superset attributes are the same.
  57. """
  58. if not isinstance(other, Subset):
  59. return NotImplemented
  60. return self.subset == other.subset and self.superset == other.superset
  61. def iterate_binary(self, k):
  62. """
  63. This is a helper function. It iterates over the
  64. binary subsets by ``k`` steps. This variable can be
  65. both positive or negative.
  66. Examples
  67. ========
  68. >>> from sympy.combinatorics import Subset
  69. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  70. >>> a.iterate_binary(-2).subset
  71. ['d']
  72. >>> a = Subset(['a', 'b', 'c'], ['a', 'b', 'c', 'd'])
  73. >>> a.iterate_binary(2).subset
  74. []
  75. See Also
  76. ========
  77. next_binary, prev_binary
  78. """
  79. bin_list = Subset.bitlist_from_subset(self.subset, self.superset)
  80. n = (int(''.join(bin_list), 2) + k) % 2**self.superset_size
  81. bits = bin(n)[2:].rjust(self.superset_size, '0')
  82. return Subset.subset_from_bitlist(self.superset, bits)
  83. def next_binary(self):
  84. """
  85. Generates the next binary ordered subset.
  86. Examples
  87. ========
  88. >>> from sympy.combinatorics import Subset
  89. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  90. >>> a.next_binary().subset
  91. ['b']
  92. >>> a = Subset(['a', 'b', 'c', 'd'], ['a', 'b', 'c', 'd'])
  93. >>> a.next_binary().subset
  94. []
  95. See Also
  96. ========
  97. prev_binary, iterate_binary
  98. """
  99. return self.iterate_binary(1)
  100. def prev_binary(self):
  101. """
  102. Generates the previous binary ordered subset.
  103. Examples
  104. ========
  105. >>> from sympy.combinatorics import Subset
  106. >>> a = Subset([], ['a', 'b', 'c', 'd'])
  107. >>> a.prev_binary().subset
  108. ['a', 'b', 'c', 'd']
  109. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  110. >>> a.prev_binary().subset
  111. ['c']
  112. See Also
  113. ========
  114. next_binary, iterate_binary
  115. """
  116. return self.iterate_binary(-1)
  117. def next_lexicographic(self):
  118. """
  119. Generates the next lexicographically ordered subset.
  120. Examples
  121. ========
  122. >>> from sympy.combinatorics import Subset
  123. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  124. >>> a.next_lexicographic().subset
  125. ['d']
  126. >>> a = Subset(['d'], ['a', 'b', 'c', 'd'])
  127. >>> a.next_lexicographic().subset
  128. []
  129. See Also
  130. ========
  131. prev_lexicographic
  132. """
  133. i = self.superset_size - 1
  134. indices = Subset.subset_indices(self.subset, self.superset)
  135. if i in indices:
  136. if i - 1 in indices:
  137. indices.remove(i - 1)
  138. else:
  139. indices.remove(i)
  140. i = i - 1
  141. while i >= 0 and i not in indices:
  142. i = i - 1
  143. if i >= 0:
  144. indices.remove(i)
  145. indices.append(i+1)
  146. else:
  147. while i not in indices and i >= 0:
  148. i = i - 1
  149. indices.append(i + 1)
  150. ret_set = []
  151. super_set = self.superset
  152. for i in indices:
  153. ret_set.append(super_set[i])
  154. return Subset(ret_set, super_set)
  155. def prev_lexicographic(self):
  156. """
  157. Generates the previous lexicographically ordered subset.
  158. Examples
  159. ========
  160. >>> from sympy.combinatorics import Subset
  161. >>> a = Subset([], ['a', 'b', 'c', 'd'])
  162. >>> a.prev_lexicographic().subset
  163. ['d']
  164. >>> a = Subset(['c','d'], ['a', 'b', 'c', 'd'])
  165. >>> a.prev_lexicographic().subset
  166. ['c']
  167. See Also
  168. ========
  169. next_lexicographic
  170. """
  171. i = self.superset_size - 1
  172. indices = Subset.subset_indices(self.subset, self.superset)
  173. while i >= 0 and i not in indices:
  174. i = i - 1
  175. if i == 0 or i - 1 in indices:
  176. indices.remove(i)
  177. else:
  178. if i >= 0:
  179. indices.remove(i)
  180. indices.append(i - 1)
  181. indices.append(self.superset_size - 1)
  182. ret_set = []
  183. super_set = self.superset
  184. for i in indices:
  185. ret_set.append(super_set[i])
  186. return Subset(ret_set, super_set)
  187. def iterate_graycode(self, k):
  188. """
  189. Helper function used for prev_gray and next_gray.
  190. It performs ``k`` step overs to get the respective Gray codes.
  191. Examples
  192. ========
  193. >>> from sympy.combinatorics import Subset
  194. >>> a = Subset([1, 2, 3], [1, 2, 3, 4])
  195. >>> a.iterate_graycode(3).subset
  196. [1, 4]
  197. >>> a.iterate_graycode(-2).subset
  198. [1, 2, 4]
  199. See Also
  200. ========
  201. next_gray, prev_gray
  202. """
  203. unranked_code = GrayCode.unrank(self.superset_size,
  204. (self.rank_gray + k) % self.cardinality)
  205. return Subset.subset_from_bitlist(self.superset,
  206. unranked_code)
  207. def next_gray(self):
  208. """
  209. Generates the next Gray code ordered subset.
  210. Examples
  211. ========
  212. >>> from sympy.combinatorics import Subset
  213. >>> a = Subset([1, 2, 3], [1, 2, 3, 4])
  214. >>> a.next_gray().subset
  215. [1, 3]
  216. See Also
  217. ========
  218. iterate_graycode, prev_gray
  219. """
  220. return self.iterate_graycode(1)
  221. def prev_gray(self):
  222. """
  223. Generates the previous Gray code ordered subset.
  224. Examples
  225. ========
  226. >>> from sympy.combinatorics import Subset
  227. >>> a = Subset([2, 3, 4], [1, 2, 3, 4, 5])
  228. >>> a.prev_gray().subset
  229. [2, 3, 4, 5]
  230. See Also
  231. ========
  232. iterate_graycode, next_gray
  233. """
  234. return self.iterate_graycode(-1)
  235. @property
  236. def rank_binary(self):
  237. """
  238. Computes the binary ordered rank.
  239. Examples
  240. ========
  241. >>> from sympy.combinatorics import Subset
  242. >>> a = Subset([], ['a','b','c','d'])
  243. >>> a.rank_binary
  244. 0
  245. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  246. >>> a.rank_binary
  247. 3
  248. See Also
  249. ========
  250. iterate_binary, unrank_binary
  251. """
  252. if self._rank_binary is None:
  253. self._rank_binary = int("".join(
  254. Subset.bitlist_from_subset(self.subset,
  255. self.superset)), 2)
  256. return self._rank_binary
  257. @property
  258. def rank_lexicographic(self):
  259. """
  260. Computes the lexicographic ranking of the subset.
  261. Examples
  262. ========
  263. >>> from sympy.combinatorics import Subset
  264. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  265. >>> a.rank_lexicographic
  266. 14
  267. >>> a = Subset([2, 4, 5], [1, 2, 3, 4, 5, 6])
  268. >>> a.rank_lexicographic
  269. 43
  270. """
  271. if self._rank_lex is None:
  272. def _ranklex(self, subset_index, i, n):
  273. if subset_index == [] or i > n:
  274. return 0
  275. if i in subset_index:
  276. subset_index.remove(i)
  277. return 1 + _ranklex(self, subset_index, i + 1, n)
  278. return 2**(n - i - 1) + _ranklex(self, subset_index, i + 1, n)
  279. indices = Subset.subset_indices(self.subset, self.superset)
  280. self._rank_lex = _ranklex(self, indices, 0, self.superset_size)
  281. return self._rank_lex
  282. @property
  283. def rank_gray(self):
  284. """
  285. Computes the Gray code ranking of the subset.
  286. Examples
  287. ========
  288. >>> from sympy.combinatorics import Subset
  289. >>> a = Subset(['c','d'], ['a','b','c','d'])
  290. >>> a.rank_gray
  291. 2
  292. >>> a = Subset([2, 4, 5], [1, 2, 3, 4, 5, 6])
  293. >>> a.rank_gray
  294. 27
  295. See Also
  296. ========
  297. iterate_graycode, unrank_gray
  298. """
  299. if self._rank_graycode is None:
  300. bits = Subset.bitlist_from_subset(self.subset, self.superset)
  301. self._rank_graycode = GrayCode(len(bits), start=bits).rank
  302. return self._rank_graycode
  303. @property
  304. def subset(self):
  305. """
  306. Gets the subset represented by the current instance.
  307. Examples
  308. ========
  309. >>> from sympy.combinatorics import Subset
  310. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  311. >>> a.subset
  312. ['c', 'd']
  313. See Also
  314. ========
  315. superset, size, superset_size, cardinality
  316. """
  317. return self._subset
  318. @property
  319. def size(self):
  320. """
  321. Gets the size of the subset.
  322. Examples
  323. ========
  324. >>> from sympy.combinatorics import Subset
  325. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  326. >>> a.size
  327. 2
  328. See Also
  329. ========
  330. subset, superset, superset_size, cardinality
  331. """
  332. return len(self.subset)
  333. @property
  334. def superset(self):
  335. """
  336. Gets the superset of the subset.
  337. Examples
  338. ========
  339. >>> from sympy.combinatorics import Subset
  340. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  341. >>> a.superset
  342. ['a', 'b', 'c', 'd']
  343. See Also
  344. ========
  345. subset, size, superset_size, cardinality
  346. """
  347. return self._superset
  348. @property
  349. def superset_size(self):
  350. """
  351. Returns the size of the superset.
  352. Examples
  353. ========
  354. >>> from sympy.combinatorics import Subset
  355. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  356. >>> a.superset_size
  357. 4
  358. See Also
  359. ========
  360. subset, superset, size, cardinality
  361. """
  362. return len(self.superset)
  363. @property
  364. def cardinality(self):
  365. """
  366. Returns the number of all possible subsets.
  367. Examples
  368. ========
  369. >>> from sympy.combinatorics import Subset
  370. >>> a = Subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  371. >>> a.cardinality
  372. 16
  373. See Also
  374. ========
  375. subset, superset, size, superset_size
  376. """
  377. return 2**(self.superset_size)
  378. @classmethod
  379. def subset_from_bitlist(self, super_set, bitlist):
  380. """
  381. Gets the subset defined by the bitlist.
  382. Examples
  383. ========
  384. >>> from sympy.combinatorics import Subset
  385. >>> Subset.subset_from_bitlist(['a', 'b', 'c', 'd'], '0011').subset
  386. ['c', 'd']
  387. See Also
  388. ========
  389. bitlist_from_subset
  390. """
  391. if len(super_set) != len(bitlist):
  392. raise ValueError("The sizes of the lists are not equal")
  393. ret_set = []
  394. for i in range(len(bitlist)):
  395. if bitlist[i] == '1':
  396. ret_set.append(super_set[i])
  397. return Subset(ret_set, super_set)
  398. @classmethod
  399. def bitlist_from_subset(self, subset, superset):
  400. """
  401. Gets the bitlist corresponding to a subset.
  402. Examples
  403. ========
  404. >>> from sympy.combinatorics import Subset
  405. >>> Subset.bitlist_from_subset(['c', 'd'], ['a', 'b', 'c', 'd'])
  406. '0011'
  407. See Also
  408. ========
  409. subset_from_bitlist
  410. """
  411. bitlist = ['0'] * len(superset)
  412. if isinstance(subset, Subset):
  413. subset = subset.subset
  414. for i in Subset.subset_indices(subset, superset):
  415. bitlist[i] = '1'
  416. return ''.join(bitlist)
  417. @classmethod
  418. def unrank_binary(self, rank, superset):
  419. """
  420. Gets the binary ordered subset of the specified rank.
  421. Examples
  422. ========
  423. >>> from sympy.combinatorics import Subset
  424. >>> Subset.unrank_binary(4, ['a', 'b', 'c', 'd']).subset
  425. ['b']
  426. See Also
  427. ========
  428. iterate_binary, rank_binary
  429. """
  430. bits = bin(rank)[2:].rjust(len(superset), '0')
  431. return Subset.subset_from_bitlist(superset, bits)
  432. @classmethod
  433. def unrank_gray(self, rank, superset):
  434. """
  435. Gets the Gray code ordered subset of the specified rank.
  436. Examples
  437. ========
  438. >>> from sympy.combinatorics import Subset
  439. >>> Subset.unrank_gray(4, ['a', 'b', 'c']).subset
  440. ['a', 'b']
  441. >>> Subset.unrank_gray(0, ['a', 'b', 'c']).subset
  442. []
  443. See Also
  444. ========
  445. iterate_graycode, rank_gray
  446. """
  447. graycode_bitlist = GrayCode.unrank(len(superset), rank)
  448. return Subset.subset_from_bitlist(superset, graycode_bitlist)
  449. @classmethod
  450. def subset_indices(self, subset, superset):
  451. """Return indices of subset in superset in a list; the list is empty
  452. if all elements of ``subset`` are not in ``superset``.
  453. Examples
  454. ========
  455. >>> from sympy.combinatorics import Subset
  456. >>> superset = [1, 3, 2, 5, 4]
  457. >>> Subset.subset_indices([3, 2, 1], superset)
  458. [1, 2, 0]
  459. >>> Subset.subset_indices([1, 6], superset)
  460. []
  461. >>> Subset.subset_indices([], superset)
  462. []
  463. """
  464. a, b = superset, subset
  465. sb = set(b)
  466. d = {}
  467. for i, ai in enumerate(a):
  468. if ai in sb:
  469. d[ai] = i
  470. sb.remove(ai)
  471. if not sb:
  472. break
  473. else:
  474. return list()
  475. return [d[bi] for bi in b]
  476. def ksubsets(superset, k):
  477. """
  478. Finds the subsets of size ``k`` in lexicographic order.
  479. This uses the itertools generator.
  480. Examples
  481. ========
  482. >>> from sympy.combinatorics.subsets import ksubsets
  483. >>> list(ksubsets([1, 2, 3], 2))
  484. [(1, 2), (1, 3), (2, 3)]
  485. >>> list(ksubsets([1, 2, 3, 4, 5], 2))
  486. [(1, 2), (1, 3), (1, 4), (1, 5), (2, 3), (2, 4), \
  487. (2, 5), (3, 4), (3, 5), (4, 5)]
  488. See Also
  489. ========
  490. Subset
  491. """
  492. return combinations(superset, k)