195 lines
4.0 KiB
Objective-C
195 lines
4.0 KiB
Objective-C
@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;
|
|
}
|