Advertisement
Guest User

Untitled

a guest
Aug 22nd, 2019
72
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 6.29 KB | None | 0 0
  1. type color bool
  2.  
  3. const (
  4. black color = true
  5. red color = false
  6. )
  7.  
  8. // StatNode of tree
  9. type StatNode struct {
  10. key int
  11. color color
  12. size int
  13. left, right *StatNode
  14. p *StatNode
  15. }
  16.  
  17. // StatTree is black and red tree
  18. type StatTree struct {
  19. null *StatNode
  20. root *StatNode
  21. }
  22.  
  23. // CreateStatTree create a tree
  24. func CreateStatTree(A []int) StatTree {
  25. r := StatTree{
  26. null: &StatNode{color: black, size: 0},
  27. }
  28. r.root = r.null
  29. for _, a := range A {
  30. r.Insert(&StatNode{key: a})
  31. }
  32. return r
  33. }
  34.  
  35. // InOrder performs inorder walk
  36. func (r *StatTree) InOrder(x *StatNode) []int {
  37. var A []int
  38. if x != r.null {
  39. A = append(A, r.InOrder(x.left)...)
  40. A = append(A, x.key)
  41. A = append(A, r.InOrder(x.right)...)
  42. }
  43. return A
  44. }
  45.  
  46. // PreOrder performs preorder walk
  47. func (r *StatTree) PreOrder(x *StatNode) []int {
  48. var A []int
  49. if x != r.null {
  50. A = append(A, x.key)
  51. A = append(A, r.InOrder(x.left)...)
  52. A = append(A, r.InOrder(x.right)...)
  53. }
  54. return A
  55. }
  56.  
  57. // PostOrder performs postorder walk
  58. func (r *StatTree) PostOrder(x *StatNode) []int {
  59. var A []int
  60. if x != r.null {
  61. A = append(A, r.InOrder(x.left)...)
  62. A = append(A, r.InOrder(x.right)...)
  63. A = append(A, x.key)
  64. }
  65. return A
  66. }
  67.  
  68. // Search for a key
  69. func (r *StatTree) Search(x *StatNode, k int) *StatNode {
  70. if x == r.null || k == x.key {
  71. return x
  72. } else if k < x.key {
  73. return r.Search(x.left, k)
  74. } else {
  75. return r.Search(x.right, k)
  76. }
  77. }
  78.  
  79. // QuickSearch for a key
  80. func (r *StatTree) QuickSearch(x *StatNode, k int) *StatNode {
  81. for x != r.null && k != x.key {
  82. if k < x.key {
  83. x = x.left
  84. } else {
  85. x = x.right
  86. }
  87. }
  88. return x
  89. }
  90.  
  91. // Min minimum StatNode
  92. func (r *StatTree) Min(x *StatNode) *StatNode {
  93. for x.left != r.null {
  94. x = x.left
  95. }
  96. return x
  97. }
  98.  
  99. // Max maximum StatNode
  100. func (r *StatTree) Max(x *StatNode) *StatNode {
  101. for x.right != r.null {
  102. x = x.right
  103. }
  104. return x
  105. }
  106.  
  107. // Successor previous StatNode
  108. func (r *StatTree) Successor(x *StatNode) *StatNode {
  109. if x.right != r.null {
  110. return r.Min(x.right)
  111. }
  112. y := x.p
  113. for y != r.null && x == y.right {
  114. x, y = y, y.p
  115. }
  116. return y
  117. }
  118.  
  119. // Preccessor post StatNode
  120. func (r *StatTree) Preccessor(x *StatNode) *StatNode {
  121. if x.left != r.null {
  122. return r.Max(x.left)
  123. }
  124. y := x.p
  125. for y != r.null && x == y.left {
  126. x, y = y, y.p
  127. }
  128. return y
  129. }
  130.  
  131. func (r *StatTree) leftRotate(x *StatNode) {
  132. y := x.right
  133. x.right = y.left
  134. if y.left != r.null {
  135. y.left.p = x
  136. }
  137. y.p = x.p
  138. if x.p == r.null {
  139. r.root = y
  140. } else if x == x.p.left {
  141. x.p.left = y
  142. } else {
  143. x.p.right = y
  144. }
  145. y.left = x
  146. x.p = y
  147. y.size = x.size
  148. x.size = x.left.size + x.right.size + 1
  149. }
  150.  
  151. func (r *StatTree) rightRotate(x *StatNode) {
  152. y := x.left
  153. x.left = y.right
  154. if y.right != r.null {
  155. y.left.p = x
  156. }
  157. y.p = x.p
  158. if x.p == r.null {
  159. r.root = y
  160. } else if x == x.p.left {
  161. x.p.left = y
  162. } else {
  163. x.p.right = y
  164. }
  165. y.right = x
  166. x.p = y
  167. y.size = x.size
  168. x.size = x.left.size + x.right.size + 1
  169. }
  170.  
  171. func (r *StatTree) insertFixUp(z *StatNode) {
  172. for z.p.color == red {
  173. if z.p == z.p.p.left {
  174. y := z.p.p.right
  175. if y.color == red {
  176. z.p.color, y.color, z.p.p.color = black, black, red
  177. z = z.p.p
  178. } else {
  179. if z == z.p.right {
  180. z = z.p
  181. r.leftRotate(z)
  182. }
  183. z.p.color, z.p.p.color = black, red
  184. r.rightRotate(z.p.p)
  185. }
  186. } else {
  187. y := z.p.p.left
  188. if y.color == red {
  189. z.p.color, y.color, z.p.p.color = black, black, red
  190. z = z.p.p
  191. } else {
  192. if z == z.p.left {
  193. z = z.p
  194. r.rightRotate(z)
  195. }
  196. z.p.color, z.p.p.color = black, red
  197. r.leftRotate(z.p.p)
  198. }
  199. }
  200. }
  201. r.root.color = black
  202. }
  203.  
  204. // Insert a StatNode
  205. func (r *StatTree) Insert(z *StatNode) {
  206. y, x := r.null, r.root
  207. for x != r.null {
  208. y = x
  209. if z.key < x.key {
  210. x = x.left
  211. } else {
  212. x = x.right
  213. }
  214. }
  215. z.p = y
  216. if y == r.null {
  217. r.root = z
  218. } else if z.key < y.key {
  219. y.left = z
  220. } else {
  221. y.right = z
  222. }
  223. z.left = r.null
  224. z.right = r.null
  225. z.color = red
  226. z.size = 1
  227. x = z.p
  228. for x != r.null {
  229. x.size++
  230. x = x.p
  231. }
  232. r.insertFixUp(z)
  233. }
  234.  
  235. func (r *StatTree) transplant(u, v *StatNode) {
  236. if u.p == r.null {
  237. r.root = v
  238. } else if u == u.p.left {
  239. u.p.left = v
  240. } else {
  241. u.p.right = v
  242. }
  243. v.p = u.p
  244. }
  245.  
  246. func (r *StatTree) deleteFixUp(x *StatNode) {
  247. for x != r.root && x.color == black {
  248. if x == x.p.left {
  249. w := x.p.right
  250. if w.color == red {
  251. w.color = black
  252. x.p.color = red
  253. r.leftRotate(x.p)
  254. w = x.p.right
  255. }
  256. if w.left.color == black && w.right.color == black {
  257. w.color = red
  258. x = x.p
  259. } else if w.right.color == black {
  260. w.left.color = black
  261. w.color = red
  262. r.rightRotate(w)
  263. w = x.p.right
  264. }
  265. w.color = x.p.color
  266. x.p.color = black
  267. w.right.color = black
  268. r.leftRotate(x.p)
  269. x = r.root
  270. } else {
  271. w := x.p.left
  272. if w.color == red {
  273. w.color = black
  274. x.p.color = red
  275. r.rightRotate(x.p)
  276. w = x.p.left
  277. }
  278. if w.right.color == black && w.left.color == black {
  279. w.color = red
  280. x = x.p
  281. } else if w.left.color == black {
  282. w.right.color = black
  283. w.color = red
  284. r.leftRotate(w)
  285. w = x.p.left
  286. }
  287. w.color = x.p.color
  288. x.p.color = black
  289. w.left.color = black
  290. r.rightRotate(x.p)
  291. x = r.root
  292. }
  293. }
  294. x.color = black
  295. }
  296.  
  297. // Delete a StatNode
  298. func (r *StatTree) Delete(z *StatNode) {
  299. y, yColor := z, z.color
  300. x := z
  301. if z.left == r.null {
  302. x = z.right
  303. r.transplant(z, z.right)
  304. } else if z.right == r.null {
  305. x = z.left
  306. r.transplant(z, z.left)
  307. } else {
  308. y = r.Min(z.right)
  309. yColor = y.color
  310. x = y.right
  311. if y.p == z {
  312. x.p = y
  313. } else {
  314. r.transplant(y, y.right)
  315. y.right = z.right
  316. y.right.p = y
  317. }
  318. r.transplant(z, y)
  319. y.left = z.left
  320. y.left.p = y
  321. y.color = z.color
  322. y.size = z.size
  323. }
  324. t := x.p
  325. for t != r.null {
  326. t.size--
  327. t = t.p
  328. }
  329. if yColor == black {
  330. r.deleteFixUp(x)
  331. }
  332. }
  333.  
  334. // Select the StatNode which is ith node
  335. func (r *StatTree) Select(x *StatNode, i int) *StatNode {
  336. s := x.left.size + 1
  337. if i == s {
  338. return x
  339. } else if i < s {
  340. return r.Select(x.left, i)
  341. } else {
  342. return r.Select(x.right, i-s)
  343. }
  344. }
  345.  
  346. // Rank the x node
  347. func (r *StatTree) Rank(x *StatNode) int {
  348. s := x.left.size + 1
  349. y := x
  350. for y != r.root {
  351. if y == y.p.right {
  352. s += y.p.left.size + 1
  353. }
  354. y = y.p
  355. }
  356. return s
  357. }
  358.  
  359. func (r *StatTree) isSizeRight(s *StatNode) bool {
  360. if s == r.null {
  361. return true
  362. }
  363. if s.left.size + s.right.size + 1 != s.size {
  364. return false
  365. }
  366. return r.isSizeRight(s.left) && r.isSizeRight(s.right)
  367. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement