RosettaCodeData/Task/AVL-tree/Objeck/avl-tree.objeck

class AVLNode {
  @key : Int;
  @balance : Int;
  @height : Int;
  @left : AVLNode;
  @right : AVLNode;
  @above : AVLNode;

  New(key : Int, above : AVLNode) {
    @key := key;
    @above := above;
  }

  method : public : GetKey() ~ Int {
    return @key;
  }

  method : public : GetLeft() ~ AVLNode {
    return @left;
  }

  method : public : GetRight() ~ AVLNode {
    return @right;
  }

  method : public : GetAbove() ~ AVLNode {
    return @above;
  }

  method : public : GetBalance() ~ Int {
    return @balance;
  }

  method : public : GetHeight() ~ Int {
    return @height;
  }

  method : public : SetBalance(balance : Int) ~ Nil {
    @balance := balance;
  }

  method : public : SetHeight(height : Int) ~ Nil {
    @height := height;
  }

  method : public : SetAbove(above : AVLNode) ~ Nil {
    @above := above;
  }

  method : public : SetLeft(left : AVLNode) ~ Nil {
    @left := left;
  }

  method : public : SetRight(right : AVLNode) ~ Nil {
    @right := right;
  }

  method : public : SetKey(key : Int) ~ Nil {
    @key := key;
  }
}

class AVLTree {
  @root : AVLNode;

  New() {}

  method : public : Insert(key : Int) ~ Bool {
    if(@root = Nil) {
      @root := AVLNode->New( key, Nil);
      return true;
    };

    n := @root;
    while(true) {
      if(n->GetKey() = key) {
        return false;
      };

      parent := n;
      goLeft := n->GetKey() > key;
      n := goLeft ? n->GetLeft() : n->GetRight();

      if(n = Nil) {
        if(goLeft) {
          parent->SetLeft(AVLNode->New( key, parent));
        } else {
          parent->SetRight(AVLNode->New( key, parent));
        };
        Rebalance(parent);
        break;
      };
    };

    return true;
  }

  method : Delete(node : AVLNode) ~ Nil {
    if (node->GetLeft() = Nil & node->GetRight() = Nil) {
      if (node ->GetAbove() = Nil) {
        @root := Nil;
      } else {
        parent := node ->GetAbove();
        if (parent->GetLeft() = node) {
          parent->SetLeft(Nil);
        } else {
          parent->SetRight(Nil);
        };
        Rebalance(parent);
      };
      return;
    };

    if (node->GetLeft() <> Nil) {
      child := node->GetLeft();
      while (child->GetRight() <> Nil) {
        child := child->GetRight();
      };
      node->SetKey(child->GetKey());
      Delete(child);
    } else {
      child := node->GetRight();
      while (child->GetLeft() <> Nil) {
        child := child->GetLeft();
      };
      node->SetKey(child->GetKey());
      Delete(child);
    };
  }

  method : public : Delete(delKey : Int) ~ Nil {
    if (@root = Nil) {
      return;
    };

    child := @root;
    while (child <> Nil) {
      node := child;
      child := delKey >= node->GetKey() ? node->GetRight() : node->GetLeft();
      if (delKey = node->GetKey()) {
        Delete(node);
        return;
      };
    };
  }

  method : Rebalance(n : AVLNode) ~ Nil {
    SetBalance(n);

    if (n->GetBalance() = -2) {
      if (Height(n->GetLeft()->GetLeft()) >= Height(n->GetLeft()->GetRight())) {
        n := RotateRight(n);
      }
      else {
        n := RotateLeftThenRight(n);
      };

    } else if (n->GetBalance() = 2) {
      if(Height(n->GetRight()->GetRight()) >= Height(n->GetRight()->GetLeft())) {
        n := RotateLeft(n);
      }
      else {
        n := RotateRightThenLeft(n);
      };
    };

    if(n->GetAbove() <> Nil) {
      Rebalance(n->GetAbove());
    } else {
      @root := n;
    };
  }

  method : RotateLeft(a : AVLNode) ~ AVLNode {
    b := a->GetRight();
    b->SetAbove(a->GetAbove());

    a->SetRight(b->GetLeft());

    if(a->GetRight() <> Nil) {
      a->GetRight()->SetAbove(a);
    };

    b->SetLeft(a);
    a->SetAbove(b);

    if (b->GetAbove() <> Nil) {
      if (b->GetAbove()->GetRight() = a) {
        b->GetAbove()->SetRight(b);
      } else {
        b->GetAbove()->SetLeft(b);
      };
    };

    SetBalance(a);
    SetBalance(b);

    return b;
  }

  method : RotateRight(a : AVLNode) ~ AVLNode {
    b := a->GetLeft();
    b->SetAbove(a->GetAbove());

    a->SetLeft(b->GetRight());

    if (a->GetLeft() <> Nil) {
      a->GetLeft()->SetAbove(a);
    };

    b->SetRight(a);
    a->SetAbove(b);

    if (b->GetAbove() <> Nil) {
      if (b->GetAbove()->GetRight() = a) {
        b->GetAbove()->SetRight(b);
      } else {
        b->GetAbove()->SetLeft(b);
      };
    };

    SetBalance(a);
    SetBalance(b);

    return b;
  }

  method : RotateLeftThenRight(n : AVLNode) ~ AVLNode {
    n->SetLeft(RotateLeft(n->GetLeft()));
    return RotateRight(n);
  }

  method : RotateRightThenLeft(n : AVLNode) ~ AVLNode {
    n->SetRight(RotateRight(n->GetRight()));
    return RotateLeft(n);
  }

  method : SetBalance(n : AVLNode) ~ Nil {
    Reheight(n);
    n->SetBalance(Height(n->GetRight()) - Height(n->GetLeft()));
  }

  method : Reheight(node : AVLNode) ~ Nil {
    if(node <> Nil) {
      node->SetHeight(1 + Int->Max(Height(node->GetLeft()), Height(node->GetRight())));
    };
  }

  method : Height(n : AVLNode) ~ Int {
    if(n = Nil) {
      return -1;
    };

    return n->GetHeight();
  }

  method : public : PrintBalance() ~ Nil {
    PrintBalance(@root);
  }

  method : PrintBalance(n : AVLNode) ~ Nil {
    if (n <> Nil) {
      PrintBalance(n->GetLeft());
      balance := n->GetBalance();
      "{$balance} "->Print();
      PrintBalance(n->GetRight());
    };
  }
}

class Test {
  function : Main(args : String[]) ~ Nil {
    tree := AVLTree->New();

    "Inserting values 1 to 10"->PrintLine();
    for(i := 1; i < 10; i+=1;) {
      tree->Insert(i);
    };

    "Printing balance: "->Print();
    tree->PrintBalance();
  }
}