gtree_redblacktree.go 25 KB

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