Initial data commit

Task/AVL-tree/Objective-C/avl-tree.m

@@ -0,0 +1,194 @@
+@implementation AVLTree
+
+-(BOOL)insertWithKey:(NSInteger)key {
+
+    if (self.root == nil) {
+        self.root = [[AVLTreeNode alloc]initWithKey:key andParent:nil];
+    } else {
+
+        AVLTreeNode *n = self.root;
+        AVLTreeNode *parent;
+
+        while (true) {
+
+            if (n.key == key) {
+                return false;
+            }
+
+            parent = n;
+
+            BOOL goLeft = n.key > key;
+            n = goLeft ? n.left : n.right;
+
+            if (n == nil) {
+
+                if (goLeft) {
+                    parent.left = [[AVLTreeNode alloc]initWithKey:key andParent:parent];
+                } else {
+                    parent.right = [[AVLTreeNode alloc]initWithKey:key andParent:parent];
+                }
+                [self rebalanceStartingAtNode:parent];
+                break;
+            }
+        }
+    }
+
+    return true;
+}
+
+-(void)rebalanceStartingAtNode:(AVLTreeNode*)n {
+
+    [self setBalance:@[n]];
+
+    if (n.balance == -2) {
+        if ([self height:(n.left.left)] >= [self height:n.left.right]) {
+            n = [self rotateRight:n];
+        } else {
+            n = [self rotateLeftThenRight:n];
+        }
+    } else if (n.balance == 2) {
+        if ([self height:n.right.right] >= [self height:n.right.left]) {
+            n = [self rotateLeft:n];
+        } else {
+            n = [self rotateRightThenLeft:n];
+        }
+    }
+
+    if (n.parent != nil) {
+        [self rebalanceStartingAtNode:n.parent];
+    } else {
+        self.root = n;
+    }
+}
+
+
+-(AVLTreeNode*)rotateRight:(AVLTreeNode*)a {
+
+    AVLTreeNode *b = a.left;
+    b.parent = a.parent;
+
+    a.left = b.right;
+
+    if (a.left != nil) {
+        a.left.parent = a;
+    }
+
+    b.right = 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;
+
+}
+
+-(AVLTreeNode*)rotateLeftThenRight:(AVLTreeNode*)n {
+
+    n.left = [self rotateLeft:n.left];
+    return [self rotateRight:n];
+
+}
+
+-(AVLTreeNode*)rotateRightThenLeft:(AVLTreeNode*)n {
+
+    n.right = [self rotateRight:n.right];
+    return [self rotateLeft:n];
+}
+
+-(AVLTreeNode*)rotateLeft:(AVLTreeNode*)a {
+
+    //set a's right node as b
+    AVLTreeNode* b = a.right;
+    //set b's parent as a's parent (which could be nil)
+    b.parent = a.parent;
+    //in case b had a left child transfer it to a
+    a.right = b.left;
+
+    // after changing a's reference to the right child, make sure the parent is set too
+    if (a.right != nil) {
+        a.right.parent = a;
+    }
+
+    // switch a over to the left to be b's left child
+    b.left = a;
+    a.parent = b;
+
+    if (b.parent != nil) {
+        if (b.parent.right == a) {
+            b.parent.right = b;
+        } else {
+            b.parent.right = b;
+        }
+    }
+
+    [self setBalance:@[a,b]];
+
+    return b;
+
+}
+
+
+
+-(void) setBalance:(NSArray*)nodesArray {
+
+    for (AVLTreeNode* n in nodesArray) {
+
+        n.balance = [self height:n.right] - [self height:n.left];
+    }
+
+}
+
+-(int)height:(AVLTreeNode*)n {
+
+    if (n == nil) {
+        return -1;
+    }
+
+    return 1 + MAX([self height:n.left], [self height:n.right]);
+}
+
+-(void)printKey:(AVLTreeNode*)n {
+    if (n != nil) {
+        [self printKey:n.left];
+        NSLog(@"%ld", n.key);
+        [self printKey:n.right];
+    }
+}
+
+-(void)printBalance:(AVLTreeNode*)n {
+    if (n != nil) {
+        [self printBalance:n.left];
+        NSLog(@"%ld", n.balance);
+        [self printBalance:n.right];
+    }
+}
+@end
+-- test
+
+int main(int argc, const char * argv[]) {
+    @autoreleasepool {
+
+        AVLTree *tree = [AVLTree new];
+        NSLog(@"inserting values 1 to 6");
+        [tree insertWithKey:1];
+        [tree insertWithKey:2];
+        [tree insertWithKey:3];
+        [tree insertWithKey:4];
+        [tree insertWithKey:5];
+        [tree insertWithKey:6];
+
+        NSLog(@"printing balance: ");
+        [tree printBalance:tree.root];
+
+        NSLog(@"printing key: ");
+        [tree printKey:tree.root];
+    }
+    return 0;
+}