Skip to content

Instantly share code, notes, and snippets.

@BobStrogg
Last active January 3, 2016 10:29
Show Gist options
  • Save BobStrogg/8449933 to your computer and use it in GitHub Desktop.
Save BobStrogg/8449933 to your computer and use it in GitHub Desktop.
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