123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182 |
- #!/usr/bin/ruby
- #
- ## https://rosettacode.org/wiki/AVL_tree
- #
- class AVLtree {
- has root = nil
- struct Node {
- Number key,
- Number balance = 0,
- Node left = nil,
- Node right = nil,
- Node parent = nil,
- }
- method insert(key) {
- if (root == nil) {
- root = Node(key)
- return true
- }
- var n = root
- var parent = nil
- loop {
- if (n.key == key) {
- return false
- }
- parent = n
- var goLeft = (n.key > key)
- n = (goLeft ? n.left : n.right)
- if (n == nil) {
- var tn = Node(key, parent: parent)
- if (goLeft) {
- parent.left = tn
- }
- else {
- parent.right = tn
- }
- self.rebalance(parent)
- break
- }
- }
- return true
- }
- method delete_key(delKey) {
- if (root == nil) { return() }
- var n = root
- var parent = root
- var delNode = nil
- var child = root
- while (child != nil) {
- parent = n
- n = child
- child = (delKey >= n.key ? n.right : n.left)
- if (delKey == n.key) {
- delNode = n
- }
- }
- if (delNode != nil) {
- delNode.key = n.key
- child = (n.left != nil ? n.left : n.right)
- if (root.key == delKey) {
- root = child
- }
- else {
- if (parent.left == n) {
- parent.left = child
- }
- else {
- parent.right = child
- }
- self.rebalance(parent)
- }
- }
- }
- method rebalance(n) {
- if (n == nil) { return() }
- self.setBalance(n)
- given (n.balance) {
- when (-2) {
- if (self.height(n.left.left) >= self.height(n.left.right)) {
- n = self.rotate(n, :right)
- }
- else {
- n = self.rotate_twice(n, :left, :right)
- }
- }
- when (2) {
- if (self.height(n.right.right) >= self.height(n.right.left)) {
- n = self.rotate(n, :left)
- }
- else {
- n = self.rotate_twice(n, :right, :left)
- }
- }
- }
- if (n.parent != nil) {
- self.rebalance(n.parent)
- }
- else {
- root = n
- }
- }
- method rotate(a, dir) {
- var b = (dir == :left ? a.right : a.left)
- b.parent = a.parent
- (dir == :left) ? (a.right = b.left)
- : (a.left = b.right)
- if (a.right != nil) {
- a.right.parent = a
- }
- b.$dir = a
- a.parent = b
- if (b.parent != nil) {
- if (b.parent.right == a) {
- b.parent.right = b
- }
- else {
- b.parent.left = b
- }
- }
- self.setBalance(a, b)
- return b
- }
- method rotate_twice(n, dir1, dir2) {
- n.left = self.rotate(n.left, dir1)
- self.rotate(n, dir2)
- }
- method height(n) {
- if (n == nil) { return -1 }
- 1 + Math.max(self.height(n.left), self.height(n.right))
- }
- method setBalance(*nodes) {
- nodes.each { |n|
- n.balance = (self.height(n.right) - self.height(n.left))
- }
- }
- method printBalance {
- self.printBalance(root)
- }
- method printBalance(n) {
- if (n != nil) {
- self.printBalance(n.left)
- print(n.balance, ' ')
- self.printBalance(n.right)
- }
- }
- }
- var tree = AVLtree()
- say "Inserting values 1 to 10"
- 10.times { |i| tree.insert(i) }
- print "Printing balance: "
- tree.printBalance
- print "\n"
|