prufer.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436
  1. from sympy.core import Basic
  2. from sympy.core.containers import Tuple
  3. from sympy.tensor.array import Array
  4. from sympy.core.sympify import _sympify
  5. from sympy.utilities.iterables import flatten, iterable
  6. from sympy.utilities.misc import as_int
  7. from collections import defaultdict
  8. class Prufer(Basic):
  9. """
  10. The Prufer correspondence is an algorithm that describes the
  11. bijection between labeled trees and the Prufer code. A Prufer
  12. code of a labeled tree is unique up to isomorphism and has
  13. a length of n - 2.
  14. Prufer sequences were first used by Heinz Prufer to give a
  15. proof of Cayley's formula.
  16. References
  17. ==========
  18. .. [1] http://mathworld.wolfram.com/LabeledTree.html
  19. """
  20. _prufer_repr = None
  21. _tree_repr = None
  22. _nodes = None
  23. _rank = None
  24. @property
  25. def prufer_repr(self):
  26. """Returns Prufer sequence for the Prufer object.
  27. This sequence is found by removing the highest numbered vertex,
  28. recording the node it was attached to, and continuing until only
  29. two vertices remain. The Prufer sequence is the list of recorded nodes.
  30. Examples
  31. ========
  32. >>> from sympy.combinatorics.prufer import Prufer
  33. >>> Prufer([[0, 3], [1, 3], [2, 3], [3, 4], [4, 5]]).prufer_repr
  34. [3, 3, 3, 4]
  35. >>> Prufer([1, 0, 0]).prufer_repr
  36. [1, 0, 0]
  37. See Also
  38. ========
  39. to_prufer
  40. """
  41. if self._prufer_repr is None:
  42. self._prufer_repr = self.to_prufer(self._tree_repr[:], self.nodes)
  43. return self._prufer_repr
  44. @property
  45. def tree_repr(self):
  46. """Returns the tree representation of the Prufer object.
  47. Examples
  48. ========
  49. >>> from sympy.combinatorics.prufer import Prufer
  50. >>> Prufer([[0, 3], [1, 3], [2, 3], [3, 4], [4, 5]]).tree_repr
  51. [[0, 3], [1, 3], [2, 3], [3, 4], [4, 5]]
  52. >>> Prufer([1, 0, 0]).tree_repr
  53. [[1, 2], [0, 1], [0, 3], [0, 4]]
  54. See Also
  55. ========
  56. to_tree
  57. """
  58. if self._tree_repr is None:
  59. self._tree_repr = self.to_tree(self._prufer_repr[:])
  60. return self._tree_repr
  61. @property
  62. def nodes(self):
  63. """Returns the number of nodes in the tree.
  64. Examples
  65. ========
  66. >>> from sympy.combinatorics.prufer import Prufer
  67. >>> Prufer([[0, 3], [1, 3], [2, 3], [3, 4], [4, 5]]).nodes
  68. 6
  69. >>> Prufer([1, 0, 0]).nodes
  70. 5
  71. """
  72. return self._nodes
  73. @property
  74. def rank(self):
  75. """Returns the rank of the Prufer sequence.
  76. Examples
  77. ========
  78. >>> from sympy.combinatorics.prufer import Prufer
  79. >>> p = Prufer([[0, 3], [1, 3], [2, 3], [3, 4], [4, 5]])
  80. >>> p.rank
  81. 778
  82. >>> p.next(1).rank
  83. 779
  84. >>> p.prev().rank
  85. 777
  86. See Also
  87. ========
  88. prufer_rank, next, prev, size
  89. """
  90. if self._rank is None:
  91. self._rank = self.prufer_rank()
  92. return self._rank
  93. @property
  94. def size(self):
  95. """Return the number of possible trees of this Prufer object.
  96. Examples
  97. ========
  98. >>> from sympy.combinatorics.prufer import Prufer
  99. >>> Prufer([0]*4).size == Prufer([6]*4).size == 1296
  100. True
  101. See Also
  102. ========
  103. prufer_rank, rank, next, prev
  104. """
  105. return self.prev(self.rank).prev().rank + 1
  106. @staticmethod
  107. def to_prufer(tree, n):
  108. """Return the Prufer sequence for a tree given as a list of edges where
  109. ``n`` is the number of nodes in the tree.
  110. Examples
  111. ========
  112. >>> from sympy.combinatorics.prufer import Prufer
  113. >>> a = Prufer([[0, 1], [0, 2], [0, 3]])
  114. >>> a.prufer_repr
  115. [0, 0]
  116. >>> Prufer.to_prufer([[0, 1], [0, 2], [0, 3]], 4)
  117. [0, 0]
  118. See Also
  119. ========
  120. prufer_repr: returns Prufer sequence of a Prufer object.
  121. """
  122. d = defaultdict(int)
  123. L = []
  124. for edge in tree:
  125. # Increment the value of the corresponding
  126. # node in the degree list as we encounter an
  127. # edge involving it.
  128. d[edge[0]] += 1
  129. d[edge[1]] += 1
  130. for i in range(n - 2):
  131. # find the smallest leaf
  132. for x in range(n):
  133. if d[x] == 1:
  134. break
  135. # find the node it was connected to
  136. y = None
  137. for edge in tree:
  138. if x == edge[0]:
  139. y = edge[1]
  140. elif x == edge[1]:
  141. y = edge[0]
  142. if y is not None:
  143. break
  144. # record and update
  145. L.append(y)
  146. for j in (x, y):
  147. d[j] -= 1
  148. if not d[j]:
  149. d.pop(j)
  150. tree.remove(edge)
  151. return L
  152. @staticmethod
  153. def to_tree(prufer):
  154. """Return the tree (as a list of edges) of the given Prufer sequence.
  155. Examples
  156. ========
  157. >>> from sympy.combinatorics.prufer import Prufer
  158. >>> a = Prufer([0, 2], 4)
  159. >>> a.tree_repr
  160. [[0, 1], [0, 2], [2, 3]]
  161. >>> Prufer.to_tree([0, 2])
  162. [[0, 1], [0, 2], [2, 3]]
  163. References
  164. ==========
  165. .. [1] https://hamberg.no/erlend/posts/2010-11-06-prufer-sequence-compact-tree-representation.html
  166. See Also
  167. ========
  168. tree_repr: returns tree representation of a Prufer object.
  169. """
  170. tree = []
  171. last = []
  172. n = len(prufer) + 2
  173. d = defaultdict(lambda: 1)
  174. for p in prufer:
  175. d[p] += 1
  176. for i in prufer:
  177. for j in range(n):
  178. # find the smallest leaf (degree = 1)
  179. if d[j] == 1:
  180. break
  181. # (i, j) is the new edge that we append to the tree
  182. # and remove from the degree dictionary
  183. d[i] -= 1
  184. d[j] -= 1
  185. tree.append(sorted([i, j]))
  186. last = [i for i in range(n) if d[i] == 1] or [0, 1]
  187. tree.append(last)
  188. return tree
  189. @staticmethod
  190. def edges(*runs):
  191. """Return a list of edges and the number of nodes from the given runs
  192. that connect nodes in an integer-labelled tree.
  193. All node numbers will be shifted so that the minimum node is 0. It is
  194. not a problem if edges are repeated in the runs; only unique edges are
  195. returned. There is no assumption made about what the range of the node
  196. labels should be, but all nodes from the smallest through the largest
  197. must be present.
  198. Examples
  199. ========
  200. >>> from sympy.combinatorics.prufer import Prufer
  201. >>> Prufer.edges([1, 2, 3], [2, 4, 5]) # a T
  202. ([[0, 1], [1, 2], [1, 3], [3, 4]], 5)
  203. Duplicate edges are removed:
  204. >>> Prufer.edges([0, 1, 2, 3], [1, 4, 5], [1, 4, 6]) # a K
  205. ([[0, 1], [1, 2], [1, 4], [2, 3], [4, 5], [4, 6]], 7)
  206. """
  207. e = set()
  208. nmin = runs[0][0]
  209. for r in runs:
  210. for i in range(len(r) - 1):
  211. a, b = r[i: i + 2]
  212. if b < a:
  213. a, b = b, a
  214. e.add((a, b))
  215. rv = []
  216. got = set()
  217. nmin = nmax = None
  218. for ei in e:
  219. for i in ei:
  220. got.add(i)
  221. nmin = min(ei[0], nmin) if nmin is not None else ei[0]
  222. nmax = max(ei[1], nmax) if nmax is not None else ei[1]
  223. rv.append(list(ei))
  224. missing = set(range(nmin, nmax + 1)) - got
  225. if missing:
  226. missing = [i + nmin for i in missing]
  227. if len(missing) == 1:
  228. msg = 'Node %s is missing.' % missing.pop()
  229. else:
  230. msg = 'Nodes %s are missing.' % list(sorted(missing))
  231. raise ValueError(msg)
  232. if nmin != 0:
  233. for i, ei in enumerate(rv):
  234. rv[i] = [n - nmin for n in ei]
  235. nmax -= nmin
  236. return sorted(rv), nmax + 1
  237. def prufer_rank(self):
  238. """Computes the rank of a Prufer sequence.
  239. Examples
  240. ========
  241. >>> from sympy.combinatorics.prufer import Prufer
  242. >>> a = Prufer([[0, 1], [0, 2], [0, 3]])
  243. >>> a.prufer_rank()
  244. 0
  245. See Also
  246. ========
  247. rank, next, prev, size
  248. """
  249. r = 0
  250. p = 1
  251. for i in range(self.nodes - 3, -1, -1):
  252. r += p*self.prufer_repr[i]
  253. p *= self.nodes
  254. return r
  255. @classmethod
  256. def unrank(self, rank, n):
  257. """Finds the unranked Prufer sequence.
  258. Examples
  259. ========
  260. >>> from sympy.combinatorics.prufer import Prufer
  261. >>> Prufer.unrank(0, 4)
  262. Prufer([0, 0])
  263. """
  264. n, rank = as_int(n), as_int(rank)
  265. L = defaultdict(int)
  266. for i in range(n - 3, -1, -1):
  267. L[i] = rank % n
  268. rank = (rank - L[i])//n
  269. return Prufer([L[i] for i in range(len(L))])
  270. def __new__(cls, *args, **kw_args):
  271. """The constructor for the Prufer object.
  272. Examples
  273. ========
  274. >>> from sympy.combinatorics.prufer import Prufer
  275. A Prufer object can be constructed from a list of edges:
  276. >>> a = Prufer([[0, 1], [0, 2], [0, 3]])
  277. >>> a.prufer_repr
  278. [0, 0]
  279. If the number of nodes is given, no checking of the nodes will
  280. be performed; it will be assumed that nodes 0 through n - 1 are
  281. present:
  282. >>> Prufer([[0, 1], [0, 2], [0, 3]], 4)
  283. Prufer([[0, 1], [0, 2], [0, 3]], 4)
  284. A Prufer object can be constructed from a Prufer sequence:
  285. >>> b = Prufer([1, 3])
  286. >>> b.tree_repr
  287. [[0, 1], [1, 3], [2, 3]]
  288. """
  289. arg0 = Array(args[0]) if args[0] else Tuple()
  290. args = (arg0,) + tuple(_sympify(arg) for arg in args[1:])
  291. ret_obj = Basic.__new__(cls, *args, **kw_args)
  292. args = [list(args[0])]
  293. if args[0] and iterable(args[0][0]):
  294. if not args[0][0]:
  295. raise ValueError(
  296. 'Prufer expects at least one edge in the tree.')
  297. if len(args) > 1:
  298. nnodes = args[1]
  299. else:
  300. nodes = set(flatten(args[0]))
  301. nnodes = max(nodes) + 1
  302. if nnodes != len(nodes):
  303. missing = set(range(nnodes)) - nodes
  304. if len(missing) == 1:
  305. msg = 'Node %s is missing.' % missing.pop()
  306. else:
  307. msg = 'Nodes %s are missing.' % list(sorted(missing))
  308. raise ValueError(msg)
  309. ret_obj._tree_repr = [list(i) for i in args[0]]
  310. ret_obj._nodes = nnodes
  311. else:
  312. ret_obj._prufer_repr = args[0]
  313. ret_obj._nodes = len(ret_obj._prufer_repr) + 2
  314. return ret_obj
  315. def next(self, delta=1):
  316. """Generates the Prufer sequence that is delta beyond the current one.
  317. Examples
  318. ========
  319. >>> from sympy.combinatorics.prufer import Prufer
  320. >>> a = Prufer([[0, 1], [0, 2], [0, 3]])
  321. >>> b = a.next(1) # == a.next()
  322. >>> b.tree_repr
  323. [[0, 2], [0, 1], [1, 3]]
  324. >>> b.rank
  325. 1
  326. See Also
  327. ========
  328. prufer_rank, rank, prev, size
  329. """
  330. return Prufer.unrank(self.rank + delta, self.nodes)
  331. def prev(self, delta=1):
  332. """Generates the Prufer sequence that is -delta before the current one.
  333. Examples
  334. ========
  335. >>> from sympy.combinatorics.prufer import Prufer
  336. >>> a = Prufer([[0, 1], [1, 2], [2, 3], [1, 4]])
  337. >>> a.rank
  338. 36
  339. >>> b = a.prev()
  340. >>> b
  341. Prufer([1, 2, 0])
  342. >>> b.rank
  343. 35
  344. See Also
  345. ========
  346. prufer_rank, rank, next, size
  347. """
  348. return Prufer.unrank(self.rank -delta, self.nodes)