TreeMap.js 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262
  1. import ArrayList from './ArrayList'
  2. import SortedMap from './SortedMap'
  3. import HashSet from './HashSet'
  4. const BLACK = 0
  5. const RED = 1
  6. function colorOf(p) {
  7. return (p == null ? BLACK : p.color)
  8. }
  9. function parentOf(p) {
  10. return (p == null ? null : p.parent)
  11. }
  12. function setColor(p, c) {
  13. if (p !== null) p.color = c
  14. }
  15. function leftOf(p) {
  16. return (p == null ? null : p.left)
  17. }
  18. function rightOf(p) {
  19. return (p == null ? null : p.right)
  20. }
  21. /**
  22. * @see http://download.oracle.com/javase/6/docs/api/java/util/TreeMap.html
  23. */
  24. export default class TreeMap extends SortedMap {
  25. constructor() {
  26. super()
  27. this.root_ = null
  28. this.size_ = 0
  29. }
  30. get(key) {
  31. let p = this.root_
  32. while (p !== null) {
  33. const cmp = key.compareTo(p.key)
  34. if (cmp < 0)
  35. p = p.left
  36. else if (cmp > 0)
  37. p = p.right
  38. else return p.value
  39. }
  40. return null
  41. }
  42. put(key, value) {
  43. if (this.root_ === null) {
  44. this.root_ = {
  45. key: key,
  46. value: value,
  47. left: null,
  48. right: null,
  49. parent: null,
  50. color: BLACK,
  51. getValue() {
  52. return this.value
  53. },
  54. getKey() {
  55. return this.key
  56. }
  57. }
  58. this.size_ = 1
  59. return null
  60. }
  61. let t = this.root_; let parent; let cmp
  62. do {
  63. parent = t
  64. cmp = key.compareTo(t.key)
  65. if (cmp < 0) {
  66. t = t.left
  67. } else if (cmp > 0) {
  68. t = t.right
  69. } else {
  70. const oldValue = t.value
  71. t.value = value
  72. return oldValue
  73. }
  74. } while (t !== null)
  75. const e = {
  76. key: key,
  77. left: null,
  78. right: null,
  79. value: value,
  80. parent: parent,
  81. color: BLACK,
  82. getValue() {
  83. return this.value
  84. },
  85. getKey() {
  86. return this.key
  87. }
  88. }
  89. if (cmp < 0)
  90. parent.left = e
  91. else parent.right = e
  92. this.fixAfterInsertion(e)
  93. this.size_++
  94. return null
  95. }
  96. /**
  97. * @param {Object} x
  98. */
  99. fixAfterInsertion(x) {
  100. let y
  101. x.color = RED
  102. while (x != null && x !== this.root_ && x.parent.color === RED)
  103. if (parentOf(x) === leftOf(parentOf(parentOf(x)))) {
  104. y = rightOf(parentOf(parentOf(x)))
  105. if (colorOf(y) === RED) {
  106. setColor(parentOf(x), BLACK)
  107. setColor(y, BLACK)
  108. setColor(parentOf(parentOf(x)), RED)
  109. x = parentOf(parentOf(x))
  110. } else {
  111. if (x === rightOf(parentOf(x))) {
  112. x = parentOf(x)
  113. this.rotateLeft(x)
  114. }
  115. setColor(parentOf(x), BLACK)
  116. setColor(parentOf(parentOf(x)), RED)
  117. this.rotateRight(parentOf(parentOf(x)))
  118. }
  119. } else {
  120. y = leftOf(parentOf(parentOf(x)))
  121. if (colorOf(y) === RED) {
  122. setColor(parentOf(x), BLACK)
  123. setColor(y, BLACK)
  124. setColor(parentOf(parentOf(x)), RED)
  125. x = parentOf(parentOf(x))
  126. } else {
  127. if (x === leftOf(parentOf(x))) {
  128. x = parentOf(x)
  129. this.rotateRight(x)
  130. }
  131. setColor(parentOf(x), BLACK)
  132. setColor(parentOf(parentOf(x)), RED)
  133. this.rotateLeft(parentOf(parentOf(x)))
  134. }
  135. }
  136. this.root_.color = BLACK
  137. }
  138. values() {
  139. const arrayList = new ArrayList()
  140. let p = this.getFirstEntry()
  141. if (p !== null) {
  142. arrayList.add(p.value)
  143. while ((p = TreeMap.successor(p)) !== null)
  144. arrayList.add(p.value)
  145. }
  146. return arrayList
  147. }
  148. entrySet() {
  149. const hashSet = new HashSet()
  150. let p = this.getFirstEntry()
  151. if (p !== null) {
  152. hashSet.add(p)
  153. while ((p = TreeMap.successor(p)) !== null)
  154. hashSet.add(p)
  155. }
  156. return hashSet
  157. }
  158. /**
  159. * @param {Object} p
  160. */
  161. rotateLeft(p) {
  162. if (p != null) {
  163. const r = p.right
  164. p.right = r.left
  165. if (r.left != null)
  166. r.left.parent = p
  167. r.parent = p.parent
  168. if (p.parent == null)
  169. this.root_ = r
  170. else if (p.parent.left === p)
  171. p.parent.left = r
  172. else
  173. p.parent.right = r
  174. r.left = p
  175. p.parent = r
  176. }
  177. }
  178. /**
  179. * @param {Object} p
  180. */
  181. rotateRight(p) {
  182. if (p != null) {
  183. const l = p.left
  184. p.left = l.right
  185. if (l.right != null)
  186. l.right.parent = p
  187. l.parent = p.parent
  188. if (p.parent == null)
  189. this.root_ = l
  190. else if (p.parent.right === p)
  191. p.parent.right = l
  192. else
  193. p.parent.left = l
  194. l.right = p
  195. p.parent = l
  196. }
  197. }
  198. /**
  199. * @return {Object}
  200. */
  201. getFirstEntry() {
  202. let p = this.root_
  203. if (p != null)
  204. while (p.left != null) p = p.left
  205. return p
  206. }
  207. /**
  208. * @param {Object} t
  209. * @return {Object}
  210. * @private
  211. */
  212. static successor(t) {
  213. let p
  214. if (t === null) {
  215. return null
  216. } else if (t.right !== null) {
  217. p = t.right
  218. while (p.left !== null)
  219. p = p.left
  220. return p
  221. } else {
  222. p = t.parent
  223. let ch = t
  224. while (p !== null && ch === p.right) {
  225. ch = p
  226. p = p.parent
  227. }
  228. return p
  229. }
  230. }
  231. size() {
  232. return this.size_
  233. }
  234. containsKey(key) {
  235. let p = this.root_
  236. while (p !== null) {
  237. const cmp = key.compareTo(p.key)
  238. if (cmp < 0)
  239. p = p.left
  240. else if (cmp > 0)
  241. p = p.right
  242. else return true
  243. }
  244. return false
  245. }
  246. }