Immutable AVL Tree in Objective-C
// | |
// AvlNode.h | |
// AVL tree / Node | |
// | |
// Created by Chris Cavanagh on 1/14/14. | |
// | |
#import <Foundation/Foundation.h> | |
typedef int (^EqualityComparer)( id v1, id v2 ); | |
@interface AvlNode : NSObject | |
+ (AvlNode *) empty; | |
- (BOOL) isEmpty; | |
- (AvlNode *) left; | |
- (AvlNode *) right; | |
- (int) count; | |
- (int) balance; | |
- (int) depth; | |
- (AvlNode *) searchNode:(id)value | |
comparer:(EqualityComparer)comparer; | |
- (AvlNode *) insertIntoNew:(int)index | |
value:(id)val; | |
- (AvlNode *) insertIntoNew:(id)val | |
comparer:(EqualityComparer)comparer; | |
- (AvlNode *) setItem:(int)index | |
value:(id)val; | |
- (AvlNode *) getNodeAt:(int)index; | |
- (AvlNode *) removeFromNew:(int)index | |
found:(BOOL *)found; | |
- (AvlNode *) removeFromNew:(id)val | |
comparer:(EqualityComparer)comparer | |
found:(BOOL *)found; | |
- (NSEnumerator *)getEnumerator:(BOOL)reverse; | |
@property (nonatomic, strong) id value; | |
@end |
// | |
// AvlNode.m | |
// AVL Tree / Node | |
// | |
// Created by Chris Cavanagh on 1/14/14. | |
// | |
#import "AvlNode.h" | |
@class AvlNode; | |
@interface NullNode : AvlNode | |
@end | |
@interface MutableAvlNode : AvlNode | |
@end | |
@interface AvlNodeEnumerator : NSEnumerator | |
- (id) initWithNode:(AvlNode *)node | |
reverse:(BOOL)reverse; | |
@end | |
@interface AvlNode () | |
{ | |
AvlNode *_left; | |
AvlNode *_right; | |
int _count; | |
int _depth; | |
} | |
@end | |
@implementation AvlNode | |
static AvlNode *_empty; | |
+ (void) initialize | |
{ | |
if ( [self class] == [AvlNode class] ) | |
{ | |
_empty = [NullNode new]; | |
} | |
} | |
+ (AvlNode *) empty | |
{ | |
return _empty; | |
} | |
- (BOOL) isEmpty | |
{ | |
return NO; | |
} | |
- (AvlNode *) left | |
{ | |
return _left; | |
} | |
- (AvlNode *) right | |
{ | |
return _right; | |
} | |
- (int) count | |
{ | |
return _count; | |
} | |
- (int) balance | |
{ | |
if ( [self isEmpty] ) return 0; | |
return _left->_depth - _right->_depth; | |
} | |
- (int) depth | |
{ | |
return _depth; | |
} | |
- (id) init | |
{ | |
if ( ( self = [super init] ) ) | |
{ | |
_right = [AvlNode empty]; | |
_left = [AvlNode empty]; | |
} | |
return self; | |
} | |
- (id) initWithValue:(id)value | |
{ | |
AvlNode *empty = [AvlNode empty]; | |
return [self initWithValue:value left:empty right:empty]; | |
} | |
- (id) initWithValue:(id)value | |
left:(AvlNode *)lt | |
right:(AvlNode *)gt | |
{ | |
if ( ( self = [super init] ) ) | |
{ | |
[self setValue:value left:lt right:gt]; | |
} | |
return self; | |
} | |
- (void) setValue:(id)value | |
left:(AvlNode *)lt | |
right:(AvlNode *)gt | |
{ | |
_value = value; | |
_left = lt; | |
_right = gt; | |
_count = 1 + _left->_count + _right->_count; | |
_depth = 1 + MAX( _left->_depth, _right->_depth ); | |
} | |
- (AvlNode *) min | |
{ | |
AvlNode *empty = [AvlNode empty]; | |
if ( [self isEmpty] ) return empty; | |
AvlNode *dict = self; | |
AvlNode *next = dict->_left; | |
while ( next != empty ) | |
{ | |
dict = next; | |
next = dict->_left; | |
} | |
return dict; | |
} | |
- (AvlNode *) fixRootBalance | |
{ | |
int bal = [self balance]; | |
if ( abs( bal ) < 2 ) return self; | |
if ( bal == 2 ) | |
{ | |
int leftBalance = [_left balance]; | |
if ( leftBalance == 1 || leftBalance == 0 ) | |
{ | |
//Easy case: | |
return [self rotateToGT]; | |
} | |
if ( leftBalance == -1 ) | |
{ | |
//Rotate LTDict: | |
AvlNode *newlt = [[self toMutableIfNecessary:_left] rotateToLT]; | |
AvlNode *newroot = [self newOrMutate:_value left:newlt right:_right]; | |
return [newroot rotateToGT]; | |
} | |
[NSException raise:@"LTDict too unbalanced" format:@"LTDict too unbalanced: %d", leftBalance]; | |
} | |
if ( bal == -2 ) | |
{ | |
int rightBalance = [_right balance]; | |
if ( rightBalance == -1 || rightBalance == 0 ) | |
{ | |
//Easy case: | |
return [self rotateToLT]; | |
} | |
if ( rightBalance == 1 ) | |
{ | |
//Rotate GTDict: | |
AvlNode *newgt = [[self toMutableIfNecessary:_right] rotateToGT]; | |
AvlNode *newroot = [self newOrMutate:_value left:_left right:newgt]; | |
return [newroot rotateToLT]; | |
} | |
[NSException raise:@"RTDict too unbalanced" format:@"RTDict too unbalanced: %d", rightBalance]; | |
} | |
//In this case we can show: |bal| > 2 | |
//if( Math.Abs(bal) > 2 ) { | |
[NSException raise:@"Tree too out of balance" format:@"Tree too out of balance: %d", bal]; | |
return nil; | |
} | |
- (AvlNode *) searchNode:(id)value | |
comparer:(EqualityComparer)comparer | |
{ | |
AvlNode *dict = self; | |
AvlNode *empty = [AvlNode empty]; | |
while ( dict != empty ) | |
{ | |
int comp = comparer( dict->_value, value ); | |
if ( comp < 0 ) | |
{ | |
dict = dict->_right; | |
} | |
else if ( comp > 0 ) | |
{ | |
dict = dict->_left; | |
} | |
else | |
{ | |
//Awesome: | |
return dict; | |
} | |
} | |
return empty; | |
} | |
/// <summary> | |
/// Return a new tree with the key-value pair inserted | |
/// If the key is already present, it replaces the value | |
/// This operation is O(Log N) where N is the number of keys | |
/// </summary> | |
- (AvlNode *) insertIntoNew:(int)index | |
value:(id)val | |
{ | |
if ( [self isEmpty] ) return [[AvlNode alloc] initWithValue:val]; | |
AvlNode *newlt = _left; | |
AvlNode *newgt = _right; | |
if ( index <= _left->_count ) | |
{ | |
newlt = [[self toMutableIfNecessary:_left] insertIntoNew:index value:val]; | |
} | |
else | |
{ | |
newgt = [[self toMutableIfNecessary:_right] insertIntoNew:index - _left->_count - 1 value:val]; | |
} | |
AvlNode *newroot = [self newOrMutate:_value left:newlt right:newgt]; | |
return [newroot fixRootBalance]; | |
} | |
/// <summary> | |
/// Return a new tree with the key-value pair inserted | |
/// If the key is already present, it replaces the value | |
/// This operation is O(Log N) where N is the number of keys | |
/// </summary> | |
- (AvlNode *) insertIntoNew:(id)val | |
comparer:(EqualityComparer)comparer | |
{ | |
if ( [self isEmpty] ) return [[AvlNode alloc] initWithValue:val]; | |
AvlNode *newlt = _left; | |
AvlNode *newgt = _right; | |
int comp = comparer( _value, val ); | |
id newv = _value; | |
if ( comp < 0 ) | |
{ | |
//Let the GTDict put it in: | |
newgt = [[self toMutableIfNecessary:_right] insertIntoNew:val comparer:comparer]; | |
} | |
else if ( comp > 0 ) | |
{ | |
//Let the LTDict put it in: | |
newlt = [[self toMutableIfNecessary:_left] insertIntoNew:val comparer:comparer]; | |
} | |
else | |
{ | |
//Replace the current value: | |
newv = val; | |
} | |
AvlNode *newroot = [self newOrMutate:newv left:newlt right:newgt]; | |
return [newroot fixRootBalance]; | |
} | |
- (AvlNode *) setItem:(int)index | |
value:(id)val | |
{ | |
AvlNode *newlt = _left; | |
AvlNode *newgt = _right; | |
if ( index < _left->_count) | |
{ | |
newlt = [[self toMutableIfNecessary:_left] setItem:index value:val]; | |
} | |
else if ( index > _left->_count ) | |
{ | |
newgt = [[self toMutableIfNecessary:_right] setItem:index - _left->_count - 1 value:val]; | |
} | |
else | |
{ | |
return [self newOrMutate:val left:newlt right:newgt]; | |
} | |
return [self newOrMutate:_value left:newlt right:newgt]; | |
} | |
- (AvlNode *) getNodeAt:(int)index | |
{ | |
if ( index < _left->_count ) return [_left getNodeAt:index]; | |
if ( index > _left->_count) return [_right getNodeAt:index - _left->_count - 1]; | |
return self; | |
} | |
/// <summary> | |
/// Try to remove the key, and return the resulting Dict | |
/// if the key is not found, old_node is Empty, else old_node is the Dict | |
/// with matching Key | |
/// </summary> | |
- (AvlNode *) removeFromNew:(int)index | |
found:(BOOL *)found | |
{ | |
if ( [self isEmpty] ) | |
{ | |
*found = NO; | |
return [AvlNode empty]; | |
} | |
if ( index < _left->_count ) | |
{ | |
AvlNode *newlt = [[self toMutableIfNecessary:_left] removeFromNew:index found:found]; | |
if ( !*found ) | |
{ | |
//Not found, so nothing changed | |
return self; | |
} | |
AvlNode *newroot = [self newOrMutate:_value left:newlt right:_right]; | |
return [newroot fixRootBalance]; | |
} | |
if ( index > _left->_count ) | |
{ | |
AvlNode *newgt = [[self toMutableIfNecessary:_right] removeFromNew:index - _left->_count - 1 found:found]; | |
if ( !*found ) | |
{ | |
//Not found, so nothing changed | |
return self; | |
} | |
AvlNode *newroot = [self newOrMutate:_value left:_left right:newgt]; | |
return [newroot fixRootBalance]; | |
} | |
//found it | |
*found = YES; | |
return [self removeRoot]; | |
} | |
/// <summary> | |
/// Try to remove the key, and return the resulting Dict | |
/// if the key is not found, old_node is Empty, else old_node is the Dict | |
/// with matching Key | |
/// </summary> | |
- (AvlNode *) removeFromNew:(id)val | |
comparer:(EqualityComparer)comparer | |
found:(BOOL *)found | |
{ | |
if ( [self isEmpty] ) | |
{ | |
*found = NO; | |
return [AvlNode empty]; | |
} | |
int comp = comparer( _value, val ); | |
if ( comp < 0 ) | |
{ | |
AvlNode *newgt = [[self toMutableIfNecessary:_right] removeFromNew:val comparer:comparer found:found]; | |
if ( !*found ) | |
{ | |
//Not found, so nothing changed | |
return self; | |
} | |
AvlNode *newroot = [self newOrMutate:_value left:_left right:newgt]; | |
return [newroot fixRootBalance]; | |
} | |
if ( comp > 0 ) | |
{ | |
AvlNode *newlt = [[self toMutableIfNecessary:_left] removeFromNew:val comparer:comparer found:found]; | |
if ( !*found ) | |
{ | |
//Not found, so nothing changed | |
return self; | |
} | |
AvlNode *newroot = [self newOrMutate:_value left:newlt right:_right]; | |
return [newroot fixRootBalance]; | |
} | |
//found it | |
*found = YES; | |
return [self removeRoot]; | |
} | |
- (AvlNode *) removeMax:(AvlNode **)max | |
{ | |
if ( [self isEmpty] ) | |
{ | |
AvlNode *empty = [AvlNode empty]; | |
*max = empty; | |
return empty; | |
} | |
if ( [_right isEmpty] ) | |
{ | |
//We are the max: | |
*max = self; | |
return _left; | |
} | |
else | |
{ | |
//Go down: | |
AvlNode *newgt = [[self toMutableIfNecessary:_right] removeMax:max]; | |
AvlNode *newroot = [self newOrMutate:_value left:_left right:newgt]; | |
return [newroot fixRootBalance]; | |
} | |
} | |
- (AvlNode *) removeMin:(AvlNode **)min | |
{ | |
if ( [self isEmpty] ) | |
{ | |
AvlNode *empty = [AvlNode empty]; | |
*min = empty; | |
return empty; | |
} | |
if ( [_left isEmpty] ) | |
{ | |
//We are the minimum: | |
*min = self; | |
return _right; | |
} | |
else | |
{ | |
//Go down: | |
AvlNode *newlt = [[self toMutableIfNecessary:_left] removeMin:min]; | |
AvlNode *newroot = [self newOrMutate:_value left:newlt right:_right]; | |
return [newroot fixRootBalance]; | |
} | |
} | |
/// <summary> | |
/// Return a new dict with the root key-value pair removed | |
/// </summary> | |
- (AvlNode *) removeRoot | |
{ | |
if ( [self isEmpty] ) | |
{ | |
return self; | |
} | |
if ( [_left isEmpty] ) | |
{ | |
return _right; | |
} | |
if ( [_right isEmpty] ) | |
{ | |
return _left; | |
} | |
//Neither are empty: | |
if ( _left->_count < _right->_count) | |
{ | |
//LTDict has fewer, so promote from GTDict to minimize depth | |
AvlNode *min; | |
AvlNode *newgt = [[self toMutableIfNecessary:_right] removeMin:&min]; | |
AvlNode *newroot = [self newOrMutate:min.value left:_left right:newgt]; | |
return [newroot fixRootBalance]; | |
} | |
else | |
{ | |
AvlNode *max; | |
AvlNode *newlt = [[self toMutableIfNecessary:_left] removeMax:&max]; | |
AvlNode *newroot = [self newOrMutate:max.value left:newlt right:_right]; | |
return [newroot fixRootBalance]; | |
} | |
} | |
/// <summary> | |
/// Move the Root into the GTDict and promote LTDict node up | |
/// If LTDict is empty, this operation returns this | |
/// </summary> | |
- (AvlNode *) rotateToGT | |
{ | |
if ( [_left isEmpty] || [self isEmpty] ) | |
{ | |
return self; | |
} | |
AvlNode *newLeft = [self toMutableIfNecessary:_left]; | |
AvlNode *lL = newLeft.left; | |
AvlNode *lR = newLeft.right; | |
AvlNode *newRight = [self newOrMutate:_value left:lR right:_right]; | |
return [newLeft newOrMutate:newLeft.value left:lL right:newRight]; | |
} | |
/// <summary> | |
/// Move the Root into the LTDict and promote GTDict node up | |
/// If GTDict is empty, this operation returns this | |
/// </summary> | |
- (AvlNode *) rotateToLT | |
{ | |
if ( [_right isEmpty] || [self isEmpty] ) | |
{ | |
return self; | |
} | |
AvlNode *newRight = [self toMutableIfNecessary:_right]; | |
AvlNode *rL = newRight.left; | |
AvlNode* rR = newRight.right; | |
AvlNode *newLeft = [self newOrMutate:_value left:_left right:rL]; | |
return [newRight newOrMutate:newRight.value left:newLeft right:rR]; | |
} | |
- (NSEnumerator *)getEnumerator:(BOOL)reverse | |
{ | |
return [[AvlNodeEnumerator alloc] initWithNode:self reverse:reverse]; | |
} | |
- (AvlNode *) toImmutable | |
{ | |
return self; | |
} | |
- (AvlNode *) newOrMutate:(id)newValue | |
left:(AvlNode *)newLeft | |
right:(AvlNode *)newRight | |
{ | |
return [[AvlNode alloc] initWithValue:newValue left:newLeft right:newRight]; | |
} | |
- (AvlNode *) toMutable | |
{ | |
return [[MutableAvlNode alloc] initWithValue:_value left:_left right:_right]; | |
} | |
- (AvlNode *) toMutableIfNecessary:(AvlNode *)node | |
{ | |
return node; | |
} | |
- (BOOL) isMutable | |
{ | |
return NO; | |
} | |
@end | |
@implementation MutableAvlNode | |
- (id) initWithValue:(id)val | |
left:(AvlNode *)lt | |
right:(AvlNode *)gt | |
{ | |
return [super initWithValue:val left:lt right:gt]; | |
} | |
- (AvlNode *) toImmutable | |
{ | |
return [[AvlNode alloc] initWithValue:self.value left:[self.left toImmutable] right:[self.right toImmutable]]; | |
} | |
- (AvlNode *) newOrMutate:(id)newValue | |
left:(AvlNode *)newLeft | |
right:(AvlNode *)newRight | |
{ | |
[self setValue:newValue left:newLeft right:newRight]; | |
return self; | |
} | |
- (AvlNode *) toMutable | |
{ | |
return self; | |
} | |
- (AvlNode *) toMutableIfNecessary:(AvlNode *)node | |
{ | |
return [node toMutable]; | |
} | |
- (BOOL) isMutable | |
{ | |
return YES; | |
} | |
@end | |
@implementation NullNode | |
- (BOOL)isEmpty | |
{ | |
return YES; | |
} | |
@end | |
@implementation AvlNodeEnumerator | |
{ | |
NSMutableArray *_toVisit; | |
BOOL _reverse; | |
BOOL _needsPop; | |
AvlNode *_root; | |
AvlNode *_this_d; | |
} | |
- (id) initWithNode:(AvlNode *)root | |
reverse:(BOOL)reverse | |
{ | |
if ( ( self = [super init] ) ) | |
{ | |
_root = root; | |
_toVisit = [NSMutableArray new]; | |
[_toVisit addObject:root]; | |
_reverse = reverse; | |
_needsPop = YES; | |
} | |
return self; | |
} | |
- (id) nextObject | |
{ | |
while ( !_needsPop || [_toVisit count] > 0 ) | |
{ | |
if ( _needsPop ) | |
{ | |
_this_d = [_toVisit lastObject]; | |
[_toVisit removeLastObject]; | |
} | |
else _needsPop = YES; | |
if ( [_this_d isEmpty] ) | |
{ | |
continue; | |
} | |
if ( _reverse ) | |
{ | |
if ( [_this_d.right isEmpty] ) | |
{ | |
//This is the next biggest value in the Dict: | |
id value = _this_d.value; | |
_this_d = _this_d.left; | |
_needsPop = NO; | |
return value; | |
} | |
else | |
{ | |
//Break it up | |
[_toVisit addObject:_this_d.left]; | |
[_toVisit addObject:[[AvlNode alloc] initWithValue:_this_d.value]]; | |
_this_d = _this_d.right; | |
_needsPop = NO; | |
} | |
} | |
else | |
{ | |
if ( [_this_d.left isEmpty] ) | |
{ | |
//This is the next biggest value in the Dict: | |
id value = _this_d.value; | |
_this_d = _this_d.right; | |
_needsPop = NO; | |
return value; | |
} | |
else | |
{ | |
//Break it up | |
if ( ![_this_d.right isEmpty] ) | |
{ | |
[_toVisit addObject:_this_d.right]; | |
} | |
[_toVisit addObject:[[AvlNode alloc] initWithValue:_this_d.value]]; | |
_this_d = _this_d.left; | |
_needsPop = NO; | |
} | |
} | |
} | |
return nil; | |
} | |
- (NSArray *)allObjects | |
{ | |
NSMutableArray *result = [NSMutableArray new]; | |
AvlNodeEnumerator *enumerator = [[AvlNodeEnumerator alloc] initWithNode:_root reverse:_reverse]; | |
id item = nil; | |
while ( nil != ( item = [enumerator nextObject] ) ) | |
{ | |
[result addObject:item]; | |
} | |
return result; | |
} | |
@end |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment