gtree_redblacktree.go 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991
  1. // Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
  2. //
  3. // This Source Code Form is subject to the terms of the MIT License.
  4. // If a copy of the MIT was not distributed with this file,
  5. // You can obtain one at https://github.com/gogf/gf.
  6. package gtree
  7. import (
  8. "bytes"
  9. "fmt"
  10. "github.com/gogf/gf/v2/container/gvar"
  11. "github.com/gogf/gf/v2/internal/json"
  12. "github.com/gogf/gf/v2/internal/rwmutex"
  13. "github.com/gogf/gf/v2/util/gconv"
  14. "github.com/gogf/gf/v2/util/gutil"
  15. )
  16. type color bool
  17. const (
  18. black, red color = true, false
  19. )
  20. // RedBlackTree holds elements of the red-black tree.
  21. type RedBlackTree struct {
  22. mu rwmutex.RWMutex
  23. root *RedBlackTreeNode
  24. size int
  25. comparator func(v1, v2 interface{}) int
  26. }
  27. // RedBlackTreeNode is a single element within the tree.
  28. type RedBlackTreeNode struct {
  29. Key interface{}
  30. Value interface{}
  31. color color
  32. left *RedBlackTreeNode
  33. right *RedBlackTreeNode
  34. parent *RedBlackTreeNode
  35. }
  36. // NewRedBlackTree instantiates a red-black tree with the custom key comparator.
  37. // The parameter `safe` is used to specify whether using tree in concurrent-safety,
  38. // which is false in default.
  39. func NewRedBlackTree(comparator func(v1, v2 interface{}) int, safe ...bool) *RedBlackTree {
  40. return &RedBlackTree{
  41. mu: rwmutex.Create(safe...),
  42. comparator: comparator,
  43. }
  44. }
  45. // NewRedBlackTreeFrom instantiates a red-black tree with the custom key comparator and `data` map.
  46. // The parameter `safe` is used to specify whether using tree in concurrent-safety,
  47. // which is false in default.
  48. func NewRedBlackTreeFrom(comparator func(v1, v2 interface{}) int, data map[interface{}]interface{}, safe ...bool) *RedBlackTree {
  49. tree := NewRedBlackTree(comparator, safe...)
  50. for k, v := range data {
  51. tree.doSet(k, v)
  52. }
  53. return tree
  54. }
  55. // SetComparator sets/changes the comparator for sorting.
  56. func (tree *RedBlackTree) SetComparator(comparator func(a, b interface{}) int) {
  57. tree.mu.Lock()
  58. defer tree.mu.Unlock()
  59. tree.comparator = comparator
  60. if tree.size > 0 {
  61. data := make(map[interface{}]interface{}, tree.size)
  62. tree.doIteratorAsc(tree.leftNode(), func(key, value interface{}) bool {
  63. data[key] = value
  64. return true
  65. })
  66. // Resort the tree if comparator is changed.
  67. tree.root = nil
  68. tree.size = 0
  69. for k, v := range data {
  70. tree.doSet(k, v)
  71. }
  72. }
  73. }
  74. // Clone returns a new tree with a copy of current tree.
  75. func (tree *RedBlackTree) Clone() *RedBlackTree {
  76. newTree := NewRedBlackTree(tree.comparator, tree.mu.IsSafe())
  77. newTree.Sets(tree.Map())
  78. return newTree
  79. }
  80. // Set inserts key-value item into the tree.
  81. func (tree *RedBlackTree) Set(key interface{}, value interface{}) {
  82. tree.mu.Lock()
  83. defer tree.mu.Unlock()
  84. tree.doSet(key, value)
  85. }
  86. // Sets batch sets key-values to the tree.
  87. func (tree *RedBlackTree) Sets(data map[interface{}]interface{}) {
  88. tree.mu.Lock()
  89. defer tree.mu.Unlock()
  90. for k, v := range data {
  91. tree.doSet(k, v)
  92. }
  93. }
  94. // doSet inserts key-value item into the tree without mutex.
  95. func (tree *RedBlackTree) doSet(key interface{}, value interface{}) {
  96. insertedNode := (*RedBlackTreeNode)(nil)
  97. if tree.root == nil {
  98. // Assert key is of comparator's type for initial tree
  99. tree.getComparator()(key, key)
  100. tree.root = &RedBlackTreeNode{Key: key, Value: value, color: red}
  101. insertedNode = tree.root
  102. } else {
  103. node := tree.root
  104. loop := true
  105. for loop {
  106. compare := tree.getComparator()(key, node.Key)
  107. switch {
  108. case compare == 0:
  109. // node.Key = key
  110. node.Value = value
  111. return
  112. case compare < 0:
  113. if node.left == nil {
  114. node.left = &RedBlackTreeNode{Key: key, Value: value, color: red}
  115. insertedNode = node.left
  116. loop = false
  117. } else {
  118. node = node.left
  119. }
  120. case compare > 0:
  121. if node.right == nil {
  122. node.right = &RedBlackTreeNode{Key: key, Value: value, color: red}
  123. insertedNode = node.right
  124. loop = false
  125. } else {
  126. node = node.right
  127. }
  128. }
  129. }
  130. insertedNode.parent = node
  131. }
  132. tree.insertCase1(insertedNode)
  133. tree.size++
  134. }
  135. // Get searches the node in the tree by `key` and returns its value or nil if key is not found in tree.
  136. func (tree *RedBlackTree) Get(key interface{}) (value interface{}) {
  137. value, _ = tree.Search(key)
  138. return
  139. }
  140. // doSetWithLockCheck checks whether value of the key exists with mutex.Lock,
  141. // if not exists, set value to the map with given `key`,
  142. // or else just return the existing value.
  143. //
  144. // When setting value, if `value` is type of <func() interface {}>,
  145. // it will be executed with mutex.Lock of the hash map,
  146. // and its return value will be set to the map with `key`.
  147. //
  148. // It returns value with given `key`.
  149. func (tree *RedBlackTree) doSetWithLockCheck(key interface{}, value interface{}) interface{} {
  150. tree.mu.Lock()
  151. defer tree.mu.Unlock()
  152. if node, found := tree.doSearch(key); found {
  153. return node.Value
  154. }
  155. if f, ok := value.(func() interface{}); ok {
  156. value = f()
  157. }
  158. if value != nil {
  159. tree.doSet(key, value)
  160. }
  161. return value
  162. }
  163. // GetOrSet returns the value by key,
  164. // or sets value with given `value` if it does not exist and then returns this value.
  165. func (tree *RedBlackTree) GetOrSet(key interface{}, value interface{}) interface{} {
  166. if v, ok := tree.Search(key); !ok {
  167. return tree.doSetWithLockCheck(key, value)
  168. } else {
  169. return v
  170. }
  171. }
  172. // GetOrSetFunc returns the value by key,
  173. // or sets value with returned value of callback function `f` if it does not exist
  174. // and then returns this value.
  175. func (tree *RedBlackTree) GetOrSetFunc(key interface{}, f func() interface{}) interface{} {
  176. if v, ok := tree.Search(key); !ok {
  177. return tree.doSetWithLockCheck(key, f())
  178. } else {
  179. return v
  180. }
  181. }
  182. // GetOrSetFuncLock returns the value by key,
  183. // or sets value with returned value of callback function `f` if it does not exist
  184. // and then returns this value.
  185. //
  186. // GetOrSetFuncLock differs with GetOrSetFunc function is that it executes function `f`
  187. // with mutex.Lock of the hash map.
  188. func (tree *RedBlackTree) GetOrSetFuncLock(key interface{}, f func() interface{}) interface{} {
  189. if v, ok := tree.Search(key); !ok {
  190. return tree.doSetWithLockCheck(key, f)
  191. } else {
  192. return v
  193. }
  194. }
  195. // GetVar returns a gvar.Var with the value by given `key`.
  196. // The returned gvar.Var is un-concurrent safe.
  197. func (tree *RedBlackTree) GetVar(key interface{}) *gvar.Var {
  198. return gvar.New(tree.Get(key))
  199. }
  200. // GetVarOrSet returns a gvar.Var with result from GetVarOrSet.
  201. // The returned gvar.Var is un-concurrent safe.
  202. func (tree *RedBlackTree) GetVarOrSet(key interface{}, value interface{}) *gvar.Var {
  203. return gvar.New(tree.GetOrSet(key, value))
  204. }
  205. // GetVarOrSetFunc returns a gvar.Var with result from GetOrSetFunc.
  206. // The returned gvar.Var is un-concurrent safe.
  207. func (tree *RedBlackTree) GetVarOrSetFunc(key interface{}, f func() interface{}) *gvar.Var {
  208. return gvar.New(tree.GetOrSetFunc(key, f))
  209. }
  210. // GetVarOrSetFuncLock returns a gvar.Var with result from GetOrSetFuncLock.
  211. // The returned gvar.Var is un-concurrent safe.
  212. func (tree *RedBlackTree) GetVarOrSetFuncLock(key interface{}, f func() interface{}) *gvar.Var {
  213. return gvar.New(tree.GetOrSetFuncLock(key, f))
  214. }
  215. // SetIfNotExist sets `value` to the map if the `key` does not exist, and then returns true.
  216. // It returns false if `key` exists, and `value` would be ignored.
  217. func (tree *RedBlackTree) SetIfNotExist(key interface{}, value interface{}) bool {
  218. if !tree.Contains(key) {
  219. tree.doSetWithLockCheck(key, value)
  220. return true
  221. }
  222. return false
  223. }
  224. // SetIfNotExistFunc sets value with return value of callback function `f`, and then returns true.
  225. // It returns false if `key` exists, and `value` would be ignored.
  226. func (tree *RedBlackTree) SetIfNotExistFunc(key interface{}, f func() interface{}) bool {
  227. if !tree.Contains(key) {
  228. tree.doSetWithLockCheck(key, f())
  229. return true
  230. }
  231. return false
  232. }
  233. // SetIfNotExistFuncLock sets value with return value of callback function `f`, and then returns true.
  234. // It returns false if `key` exists, and `value` would be ignored.
  235. //
  236. // SetIfNotExistFuncLock differs with SetIfNotExistFunc function is that
  237. // it executes function `f` with mutex.Lock of the hash map.
  238. func (tree *RedBlackTree) SetIfNotExistFuncLock(key interface{}, f func() interface{}) bool {
  239. if !tree.Contains(key) {
  240. tree.doSetWithLockCheck(key, f)
  241. return true
  242. }
  243. return false
  244. }
  245. // Contains checks whether `key` exists in the tree.
  246. func (tree *RedBlackTree) Contains(key interface{}) bool {
  247. _, ok := tree.Search(key)
  248. return ok
  249. }
  250. // doRemove removes the node from the tree by `key` without mutex.
  251. func (tree *RedBlackTree) doRemove(key interface{}) (value interface{}) {
  252. child := (*RedBlackTreeNode)(nil)
  253. node, found := tree.doSearch(key)
  254. if !found {
  255. return
  256. }
  257. value = node.Value
  258. if node.left != nil && node.right != nil {
  259. p := node.left.maximumNode()
  260. node.Key = p.Key
  261. node.Value = p.Value
  262. node = p
  263. }
  264. if node.left == nil || node.right == nil {
  265. if node.right == nil {
  266. child = node.left
  267. } else {
  268. child = node.right
  269. }
  270. if node.color == black {
  271. node.color = tree.nodeColor(child)
  272. tree.deleteCase1(node)
  273. }
  274. tree.replaceNode(node, child)
  275. if node.parent == nil && child != nil {
  276. child.color = black
  277. }
  278. }
  279. tree.size--
  280. return
  281. }
  282. // Remove removes the node from the tree by `key`.
  283. func (tree *RedBlackTree) Remove(key interface{}) (value interface{}) {
  284. tree.mu.Lock()
  285. defer tree.mu.Unlock()
  286. return tree.doRemove(key)
  287. }
  288. // Removes batch deletes values of the tree by `keys`.
  289. func (tree *RedBlackTree) Removes(keys []interface{}) {
  290. tree.mu.Lock()
  291. defer tree.mu.Unlock()
  292. for _, key := range keys {
  293. tree.doRemove(key)
  294. }
  295. }
  296. // IsEmpty returns true if tree does not contain any nodes.
  297. func (tree *RedBlackTree) IsEmpty() bool {
  298. return tree.Size() == 0
  299. }
  300. // Size returns number of nodes in the tree.
  301. func (tree *RedBlackTree) Size() int {
  302. tree.mu.RLock()
  303. defer tree.mu.RUnlock()
  304. return tree.size
  305. }
  306. // Keys returns all keys in asc order.
  307. func (tree *RedBlackTree) Keys() []interface{} {
  308. var (
  309. keys = make([]interface{}, tree.Size())
  310. index = 0
  311. )
  312. tree.IteratorAsc(func(key, value interface{}) bool {
  313. keys[index] = key
  314. index++
  315. return true
  316. })
  317. return keys
  318. }
  319. // Values returns all values in asc order based on the key.
  320. func (tree *RedBlackTree) Values() []interface{} {
  321. var (
  322. values = make([]interface{}, tree.Size())
  323. index = 0
  324. )
  325. tree.IteratorAsc(func(key, value interface{}) bool {
  326. values[index] = value
  327. index++
  328. return true
  329. })
  330. return values
  331. }
  332. // Map returns all key-value items as map.
  333. func (tree *RedBlackTree) Map() map[interface{}]interface{} {
  334. m := make(map[interface{}]interface{}, tree.Size())
  335. tree.IteratorAsc(func(key, value interface{}) bool {
  336. m[key] = value
  337. return true
  338. })
  339. return m
  340. }
  341. // MapStrAny returns all key-value items as map[string]interface{}.
  342. func (tree *RedBlackTree) MapStrAny() map[string]interface{} {
  343. m := make(map[string]interface{}, tree.Size())
  344. tree.IteratorAsc(func(key, value interface{}) bool {
  345. m[gconv.String(key)] = value
  346. return true
  347. })
  348. return m
  349. }
  350. // Left returns the left-most (min) node or nil if tree is empty.
  351. func (tree *RedBlackTree) Left() *RedBlackTreeNode {
  352. tree.mu.RLock()
  353. defer tree.mu.RUnlock()
  354. node := tree.leftNode()
  355. if tree.mu.IsSafe() {
  356. return &RedBlackTreeNode{
  357. Key: node.Key,
  358. Value: node.Value,
  359. }
  360. }
  361. return node
  362. }
  363. // Right returns the right-most (max) node or nil if tree is empty.
  364. func (tree *RedBlackTree) Right() *RedBlackTreeNode {
  365. tree.mu.RLock()
  366. defer tree.mu.RUnlock()
  367. node := tree.rightNode()
  368. if tree.mu.IsSafe() {
  369. return &RedBlackTreeNode{
  370. Key: node.Key,
  371. Value: node.Value,
  372. }
  373. }
  374. return node
  375. }
  376. // leftNode returns the left-most (min) node or nil if tree is empty.
  377. func (tree *RedBlackTree) leftNode() *RedBlackTreeNode {
  378. p := (*RedBlackTreeNode)(nil)
  379. n := tree.root
  380. for n != nil {
  381. p = n
  382. n = n.left
  383. }
  384. return p
  385. }
  386. // rightNode returns the right-most (max) node or nil if tree is empty.
  387. func (tree *RedBlackTree) rightNode() *RedBlackTreeNode {
  388. p := (*RedBlackTreeNode)(nil)
  389. n := tree.root
  390. for n != nil {
  391. p = n
  392. n = n.right
  393. }
  394. return p
  395. }
  396. // Floor Finds floor node of the input key, return the floor node or nil if no floor node is found.
  397. // Second return parameter is true if floor was found, otherwise false.
  398. //
  399. // Floor node is defined as the largest node that its key is smaller than or equal to the given `key`.
  400. // A floor node may not be found, either because the tree is empty, or because
  401. // all nodes in the tree are larger than the given node.
  402. func (tree *RedBlackTree) Floor(key interface{}) (floor *RedBlackTreeNode, found bool) {
  403. tree.mu.RLock()
  404. defer tree.mu.RUnlock()
  405. n := tree.root
  406. for n != nil {
  407. compare := tree.getComparator()(key, n.Key)
  408. switch {
  409. case compare == 0:
  410. return n, true
  411. case compare < 0:
  412. n = n.left
  413. case compare > 0:
  414. floor, found = n, true
  415. n = n.right
  416. }
  417. }
  418. if found {
  419. return
  420. }
  421. return nil, false
  422. }
  423. // Ceiling finds ceiling node of the input key, return the ceiling node or nil if no ceiling node is found.
  424. // Second return parameter is true if ceiling was found, otherwise false.
  425. //
  426. // Ceiling node is defined as the smallest node that its key is larger than or equal to the given `key`.
  427. // A ceiling node may not be found, either because the tree is empty, or because
  428. // all nodes in the tree are smaller than the given node.
  429. func (tree *RedBlackTree) Ceiling(key interface{}) (ceiling *RedBlackTreeNode, found bool) {
  430. tree.mu.RLock()
  431. defer tree.mu.RUnlock()
  432. n := tree.root
  433. for n != nil {
  434. compare := tree.getComparator()(key, n.Key)
  435. switch {
  436. case compare == 0:
  437. return n, true
  438. case compare > 0:
  439. n = n.right
  440. case compare < 0:
  441. ceiling, found = n, true
  442. n = n.left
  443. }
  444. }
  445. if found {
  446. return
  447. }
  448. return nil, false
  449. }
  450. // Iterator is alias of IteratorAsc.
  451. func (tree *RedBlackTree) Iterator(f func(key, value interface{}) bool) {
  452. tree.IteratorAsc(f)
  453. }
  454. // IteratorFrom is alias of IteratorAscFrom.
  455. func (tree *RedBlackTree) IteratorFrom(key interface{}, match bool, f func(key, value interface{}) bool) {
  456. tree.IteratorAscFrom(key, match, f)
  457. }
  458. // IteratorAsc iterates the tree readonly in ascending order with given callback function `f`.
  459. // If `f` returns true, then it continues iterating; or false to stop.
  460. func (tree *RedBlackTree) IteratorAsc(f func(key, value interface{}) bool) {
  461. tree.mu.RLock()
  462. defer tree.mu.RUnlock()
  463. tree.doIteratorAsc(tree.leftNode(), f)
  464. }
  465. // IteratorAscFrom iterates the tree readonly in ascending order with given callback function `f`.
  466. // The parameter `key` specifies the start entry for iterating. The `match` specifies whether
  467. // starting iterating if the `key` is fully matched, or else using index searching iterating.
  468. // If `f` returns true, then it continues iterating; or false to stop.
  469. func (tree *RedBlackTree) IteratorAscFrom(key interface{}, match bool, f func(key, value interface{}) bool) {
  470. tree.mu.RLock()
  471. defer tree.mu.RUnlock()
  472. node, found := tree.doSearch(key)
  473. if match {
  474. if found {
  475. tree.doIteratorAsc(node, f)
  476. }
  477. } else {
  478. tree.doIteratorAsc(node, f)
  479. }
  480. }
  481. func (tree *RedBlackTree) doIteratorAsc(node *RedBlackTreeNode, f func(key, value interface{}) bool) {
  482. loop:
  483. if node == nil {
  484. return
  485. }
  486. if !f(node.Key, node.Value) {
  487. return
  488. }
  489. if node.right != nil {
  490. node = node.right
  491. for node.left != nil {
  492. node = node.left
  493. }
  494. goto loop
  495. }
  496. if node.parent != nil {
  497. old := node
  498. for node.parent != nil {
  499. node = node.parent
  500. if tree.getComparator()(old.Key, node.Key) <= 0 {
  501. goto loop
  502. }
  503. }
  504. }
  505. }
  506. // IteratorDesc iterates the tree readonly in descending order with given callback function `f`.
  507. // If `f` returns true, then it continues iterating; or false to stop.
  508. func (tree *RedBlackTree) IteratorDesc(f func(key, value interface{}) bool) {
  509. tree.mu.RLock()
  510. defer tree.mu.RUnlock()
  511. tree.doIteratorDesc(tree.rightNode(), f)
  512. }
  513. // IteratorDescFrom iterates the tree readonly in descending order with given callback function `f`.
  514. // The parameter `key` specifies the start entry for iterating. The `match` specifies whether
  515. // starting iterating if the `key` is fully matched, or else using index searching iterating.
  516. // If `f` returns true, then it continues iterating; or false to stop.
  517. func (tree *RedBlackTree) IteratorDescFrom(key interface{}, match bool, f func(key, value interface{}) bool) {
  518. tree.mu.RLock()
  519. defer tree.mu.RUnlock()
  520. node, found := tree.doSearch(key)
  521. if match {
  522. if found {
  523. tree.doIteratorDesc(node, f)
  524. }
  525. } else {
  526. tree.doIteratorDesc(node, f)
  527. }
  528. }
  529. func (tree *RedBlackTree) doIteratorDesc(node *RedBlackTreeNode, f func(key, value interface{}) bool) {
  530. loop:
  531. if node == nil {
  532. return
  533. }
  534. if !f(node.Key, node.Value) {
  535. return
  536. }
  537. if node.left != nil {
  538. node = node.left
  539. for node.right != nil {
  540. node = node.right
  541. }
  542. goto loop
  543. }
  544. if node.parent != nil {
  545. old := node
  546. for node.parent != nil {
  547. node = node.parent
  548. if tree.getComparator()(old.Key, node.Key) >= 0 {
  549. goto loop
  550. }
  551. }
  552. }
  553. }
  554. // Clear removes all nodes from the tree.
  555. func (tree *RedBlackTree) Clear() {
  556. tree.mu.Lock()
  557. defer tree.mu.Unlock()
  558. tree.root = nil
  559. tree.size = 0
  560. }
  561. // Replace the data of the tree with given `data`.
  562. func (tree *RedBlackTree) Replace(data map[interface{}]interface{}) {
  563. tree.mu.Lock()
  564. defer tree.mu.Unlock()
  565. tree.root = nil
  566. tree.size = 0
  567. for k, v := range data {
  568. tree.doSet(k, v)
  569. }
  570. }
  571. // String returns a string representation of container.
  572. func (tree *RedBlackTree) String() string {
  573. if tree == nil {
  574. return ""
  575. }
  576. tree.mu.RLock()
  577. defer tree.mu.RUnlock()
  578. str := ""
  579. if tree.size != 0 {
  580. tree.output(tree.root, "", true, &str)
  581. }
  582. return str
  583. }
  584. // Print prints the tree to stdout.
  585. func (tree *RedBlackTree) Print() {
  586. fmt.Println(tree.String())
  587. }
  588. // Search searches the tree with given `key`.
  589. // Second return parameter `found` is true if key was found, otherwise false.
  590. func (tree *RedBlackTree) Search(key interface{}) (value interface{}, found bool) {
  591. tree.mu.RLock()
  592. defer tree.mu.RUnlock()
  593. node, found := tree.doSearch(key)
  594. if found {
  595. return node.Value, true
  596. }
  597. return nil, false
  598. }
  599. // Flip exchanges key-value of the tree to value-key.
  600. // Note that you should guarantee the value is the same type as key,
  601. // or else the comparator would panic.
  602. //
  603. // If the type of value is different with key, you pass the new `comparator`.
  604. func (tree *RedBlackTree) Flip(comparator ...func(v1, v2 interface{}) int) {
  605. t := (*RedBlackTree)(nil)
  606. if len(comparator) > 0 {
  607. t = NewRedBlackTree(comparator[0], tree.mu.IsSafe())
  608. } else {
  609. t = NewRedBlackTree(tree.comparator, tree.mu.IsSafe())
  610. }
  611. tree.IteratorAsc(func(key, value interface{}) bool {
  612. t.doSet(value, key)
  613. return true
  614. })
  615. tree.mu.Lock()
  616. tree.root = t.root
  617. tree.size = t.size
  618. tree.mu.Unlock()
  619. }
  620. func (tree *RedBlackTree) output(node *RedBlackTreeNode, prefix string, isTail bool, str *string) {
  621. if node.right != nil {
  622. newPrefix := prefix
  623. if isTail {
  624. newPrefix += "│ "
  625. } else {
  626. newPrefix += " "
  627. }
  628. tree.output(node.right, newPrefix, false, str)
  629. }
  630. *str += prefix
  631. if isTail {
  632. *str += "└── "
  633. } else {
  634. *str += "┌── "
  635. }
  636. *str += fmt.Sprintf("%v\n", node.Key)
  637. if node.left != nil {
  638. newPrefix := prefix
  639. if isTail {
  640. newPrefix += " "
  641. } else {
  642. newPrefix += "│ "
  643. }
  644. tree.output(node.left, newPrefix, true, str)
  645. }
  646. }
  647. // doSearch searches the tree with given `key` without mutex.
  648. // It returns the node if found or otherwise nil.
  649. func (tree *RedBlackTree) doSearch(key interface{}) (node *RedBlackTreeNode, found bool) {
  650. node = tree.root
  651. for node != nil {
  652. compare := tree.getComparator()(key, node.Key)
  653. switch {
  654. case compare == 0:
  655. return node, true
  656. case compare < 0:
  657. node = node.left
  658. case compare > 0:
  659. node = node.right
  660. }
  661. }
  662. return node, false
  663. }
  664. func (node *RedBlackTreeNode) grandparent() *RedBlackTreeNode {
  665. if node != nil && node.parent != nil {
  666. return node.parent.parent
  667. }
  668. return nil
  669. }
  670. func (node *RedBlackTreeNode) uncle() *RedBlackTreeNode {
  671. if node == nil || node.parent == nil || node.parent.parent == nil {
  672. return nil
  673. }
  674. return node.parent.sibling()
  675. }
  676. func (node *RedBlackTreeNode) sibling() *RedBlackTreeNode {
  677. if node == nil || node.parent == nil {
  678. return nil
  679. }
  680. if node == node.parent.left {
  681. return node.parent.right
  682. }
  683. return node.parent.left
  684. }
  685. func (tree *RedBlackTree) rotateLeft(node *RedBlackTreeNode) {
  686. right := node.right
  687. tree.replaceNode(node, right)
  688. node.right = right.left
  689. if right.left != nil {
  690. right.left.parent = node
  691. }
  692. right.left = node
  693. node.parent = right
  694. }
  695. func (tree *RedBlackTree) rotateRight(node *RedBlackTreeNode) {
  696. left := node.left
  697. tree.replaceNode(node, left)
  698. node.left = left.right
  699. if left.right != nil {
  700. left.right.parent = node
  701. }
  702. left.right = node
  703. node.parent = left
  704. }
  705. func (tree *RedBlackTree) replaceNode(old *RedBlackTreeNode, new *RedBlackTreeNode) {
  706. if old.parent == nil {
  707. tree.root = new
  708. } else {
  709. if old == old.parent.left {
  710. old.parent.left = new
  711. } else {
  712. old.parent.right = new
  713. }
  714. }
  715. if new != nil {
  716. new.parent = old.parent
  717. }
  718. }
  719. func (tree *RedBlackTree) insertCase1(node *RedBlackTreeNode) {
  720. if node.parent == nil {
  721. node.color = black
  722. } else {
  723. tree.insertCase2(node)
  724. }
  725. }
  726. func (tree *RedBlackTree) insertCase2(node *RedBlackTreeNode) {
  727. if tree.nodeColor(node.parent) == black {
  728. return
  729. }
  730. tree.insertCase3(node)
  731. }
  732. func (tree *RedBlackTree) insertCase3(node *RedBlackTreeNode) {
  733. uncle := node.uncle()
  734. if tree.nodeColor(uncle) == red {
  735. node.parent.color = black
  736. uncle.color = black
  737. node.grandparent().color = red
  738. tree.insertCase1(node.grandparent())
  739. } else {
  740. tree.insertCase4(node)
  741. }
  742. }
  743. func (tree *RedBlackTree) insertCase4(node *RedBlackTreeNode) {
  744. grandparent := node.grandparent()
  745. if node == node.parent.right && node.parent == grandparent.left {
  746. tree.rotateLeft(node.parent)
  747. node = node.left
  748. } else if node == node.parent.left && node.parent == grandparent.right {
  749. tree.rotateRight(node.parent)
  750. node = node.right
  751. }
  752. tree.insertCase5(node)
  753. }
  754. func (tree *RedBlackTree) insertCase5(node *RedBlackTreeNode) {
  755. node.parent.color = black
  756. grandparent := node.grandparent()
  757. grandparent.color = red
  758. if node == node.parent.left && node.parent == grandparent.left {
  759. tree.rotateRight(grandparent)
  760. } else if node == node.parent.right && node.parent == grandparent.right {
  761. tree.rotateLeft(grandparent)
  762. }
  763. }
  764. func (node *RedBlackTreeNode) maximumNode() *RedBlackTreeNode {
  765. if node == nil {
  766. return nil
  767. }
  768. for node.right != nil {
  769. return node.right
  770. }
  771. return node
  772. }
  773. func (tree *RedBlackTree) deleteCase1(node *RedBlackTreeNode) {
  774. if node.parent == nil {
  775. return
  776. }
  777. tree.deleteCase2(node)
  778. }
  779. func (tree *RedBlackTree) deleteCase2(node *RedBlackTreeNode) {
  780. sibling := node.sibling()
  781. if tree.nodeColor(sibling) == red {
  782. node.parent.color = red
  783. sibling.color = black
  784. if node == node.parent.left {
  785. tree.rotateLeft(node.parent)
  786. } else {
  787. tree.rotateRight(node.parent)
  788. }
  789. }
  790. tree.deleteCase3(node)
  791. }
  792. func (tree *RedBlackTree) deleteCase3(node *RedBlackTreeNode) {
  793. sibling := node.sibling()
  794. if tree.nodeColor(node.parent) == black &&
  795. tree.nodeColor(sibling) == black &&
  796. tree.nodeColor(sibling.left) == black &&
  797. tree.nodeColor(sibling.right) == black {
  798. sibling.color = red
  799. tree.deleteCase1(node.parent)
  800. } else {
  801. tree.deleteCase4(node)
  802. }
  803. }
  804. func (tree *RedBlackTree) deleteCase4(node *RedBlackTreeNode) {
  805. sibling := node.sibling()
  806. if tree.nodeColor(node.parent) == red &&
  807. tree.nodeColor(sibling) == black &&
  808. tree.nodeColor(sibling.left) == black &&
  809. tree.nodeColor(sibling.right) == black {
  810. sibling.color = red
  811. node.parent.color = black
  812. } else {
  813. tree.deleteCase5(node)
  814. }
  815. }
  816. func (tree *RedBlackTree) deleteCase5(node *RedBlackTreeNode) {
  817. sibling := node.sibling()
  818. if node == node.parent.left &&
  819. tree.nodeColor(sibling) == black &&
  820. tree.nodeColor(sibling.left) == red &&
  821. tree.nodeColor(sibling.right) == black {
  822. sibling.color = red
  823. sibling.left.color = black
  824. tree.rotateRight(sibling)
  825. } else if node == node.parent.right &&
  826. tree.nodeColor(sibling) == black &&
  827. tree.nodeColor(sibling.right) == red &&
  828. tree.nodeColor(sibling.left) == black {
  829. sibling.color = red
  830. sibling.right.color = black
  831. tree.rotateLeft(sibling)
  832. }
  833. tree.deleteCase6(node)
  834. }
  835. func (tree *RedBlackTree) deleteCase6(node *RedBlackTreeNode) {
  836. sibling := node.sibling()
  837. sibling.color = tree.nodeColor(node.parent)
  838. node.parent.color = black
  839. if node == node.parent.left && tree.nodeColor(sibling.right) == red {
  840. sibling.right.color = black
  841. tree.rotateLeft(node.parent)
  842. } else if tree.nodeColor(sibling.left) == red {
  843. sibling.left.color = black
  844. tree.rotateRight(node.parent)
  845. }
  846. }
  847. func (tree *RedBlackTree) nodeColor(node *RedBlackTreeNode) color {
  848. if node == nil {
  849. return black
  850. }
  851. return node.color
  852. }
  853. // MarshalJSON implements the interface MarshalJSON for json.Marshal.
  854. func (tree RedBlackTree) MarshalJSON() (jsonBytes []byte, err error) {
  855. if tree.root == nil {
  856. return []byte("null"), nil
  857. }
  858. buffer := bytes.NewBuffer(nil)
  859. buffer.WriteByte('{')
  860. tree.Iterator(func(key, value interface{}) bool {
  861. valueBytes, valueJsonErr := json.Marshal(value)
  862. if valueJsonErr != nil {
  863. err = valueJsonErr
  864. return false
  865. }
  866. if buffer.Len() > 1 {
  867. buffer.WriteByte(',')
  868. }
  869. buffer.WriteString(fmt.Sprintf(`"%v":%s`, key, valueBytes))
  870. return true
  871. })
  872. buffer.WriteByte('}')
  873. return buffer.Bytes(), nil
  874. }
  875. // UnmarshalJSON implements the interface UnmarshalJSON for json.Unmarshal.
  876. func (tree *RedBlackTree) UnmarshalJSON(b []byte) error {
  877. tree.mu.Lock()
  878. defer tree.mu.Unlock()
  879. if tree.comparator == nil {
  880. tree.comparator = gutil.ComparatorString
  881. }
  882. var data map[string]interface{}
  883. if err := json.UnmarshalUseNumber(b, &data); err != nil {
  884. return err
  885. }
  886. for k, v := range data {
  887. tree.doSet(k, v)
  888. }
  889. return nil
  890. }
  891. // UnmarshalValue is an interface implement which sets any type of value for map.
  892. func (tree *RedBlackTree) UnmarshalValue(value interface{}) (err error) {
  893. tree.mu.Lock()
  894. defer tree.mu.Unlock()
  895. if tree.comparator == nil {
  896. tree.comparator = gutil.ComparatorString
  897. }
  898. for k, v := range gconv.Map(value) {
  899. tree.doSet(k, v)
  900. }
  901. return
  902. }
  903. // getComparator returns the comparator if it's previously set,
  904. // or else it panics.
  905. func (tree *RedBlackTree) getComparator() func(a, b interface{}) int {
  906. if tree.comparator == nil {
  907. panic("comparator is missing for tree")
  908. }
  909. return tree.comparator
  910. }