@interface LCTrie : NSObject


/// 元素数量（key or value 的数量）
-(int)size;


/// 是否为空
-(BOOL)isEmpty;


/// 清空
-(void)clear;


/// 是否包含key
/// @param key key（可以理解为一个单词）
-(BOOL)contains:(NSString*)key;


/// 获取key对应的value
/// @param key key
-(id)get:(NSString*)key;


/// 添加或更改一对key:value元素,返回旧value
/// @param key key (可以理解为一个单词)
/// @param value value（可以理解为key的附属信息，可以为nil.这是一个字典树，key的信息相对重要）
-(id)add:(NSString*)key value:(id)value;


/// 移除元素
/// @param key key
-(id)remove:(NSString*)key;


/// 是否包含前缀
/// @param prefix 前缀
-(BOOL)starsWith:(NSString*)prefix;



@end

@interface LCTrieNode : NSObject
@property(nonatomic,strong)LCTrieNode * parent;
@property(nonatomic,strong)NSMutableDictionary* children;
@property(nonatomic,assign)NSString* character;
@property(nonatomic,strong)id value;
@property(nonatomic,assign)BOOL isWord;
+(instancetype)trieNodeWithParent:(LCTrieNode*)parent;
@end

@implementation LCTrieNode

+(instancetype)trieNodeWithParent:(LCTrieNode*)parent{
    LCTrieNode * node = [[LCTrieNode alloc] init];
    node.parent = parent;
    return node;
}

@end



@interface LCTrie()
@property(nonatomic,assign)int          size;
@property(nonatomic,strong)LCTrieNode * root;
@end


@implementation LCTrie


/// 元素数量
-(int)size{
    return _size;
}



/// 是否为空
-(BOOL)isEmpty{
    return _size == 0;
    
}


/// 清空
-(void)clear{
    _size = 0;
    _root = nil;
}


/// 是否包含key
/// @param key key
-(BOOL)contains:(NSString*)key{
    LCTrieNode * node = [self __nodeWithKey:key];
    return node && node.isWord;
}


/// 获取key对应的value
/// @param key key
-(id)get:(NSString*)key{
    LCTrieNode * node = [self __nodeWithKey:key];
    return (node&&node.isWord)?node.value:nil;
}


/// 添加元素
/// @param key key
/// @param value value
-(id)add:(NSString*)key value:(id)value{
    [self __keyCheck:key];
    
    if (self.root == nil) {
        self.root = [LCTrieNode trieNodeWithParent:nil];
    }
    
    LCTrieNode * node = _root;
    
    NSUInteger len = key.length;
    for (int i = 0; i<len; i++) {
        
        
        //node是否存在子节点，如果children不存在，创建children;
        if (node.children == nil) {
            node.children = [NSMutableDictionary dictionary];
        }
        
        //子节点内是否存在该字符对应的node,不存在则创建childNode;
        NSString* c = [key substringWithRange:NSMakeRange(i, 1)];
        LCTrieNode * childNode = [node.children objectForKey:c];
        if (childNode == nil) {
            childNode = [LCTrieNode trieNodeWithParent:node];
            childNode.character = c;
            [node.children setObject:childNode forKey:c];
        }
        
        node = childNode;
    }
    
    if (node.isWord) {
        id oldValue = node.value;
        node.value = value;
        return oldValue;
    }
    
    node.isWord = true;
    node.value = value;
    self.size++;
    return nil;
}


/// 移除元素
/// @param key key
-(id)remove:(NSString*)key{
    //找到最后一个节点
    LCTrieNode * node = [self __nodeWithKey:key];
    
    //1、如果节点不存在或key不是一个单词结尾 不用处理
    if(node == nil || !node.isWord ) return nil;
    
    id oldValue = node.value;
    
    //2、如果还有子节点
    if (node.children && node.children.count) {
        node.isWord = false;
        node.value = nil;
        return oldValue;
    }
    
    //3、如果没有子节点
    LCTrieNode * parent = nil;
    while ((parent = node.parent)) {
        [parent.children removeObjectForKey:node.character];
        if(parent.isWord || parent.children.count) break;
        node = parent;
    }
    
    return oldValue;
}


/// 是否包含前缀
/// @param prefix 前缀
-(BOOL)starsWith:(NSString*)prefix{
    return  [self __nodeWithKey:prefix]!=nil;
}


#pragma mark - 内部方法

-(void)__keyCheck:(NSString*)key{
    if (key ==nil && key.length == 0) {
        NSException * exception = [NSException exceptionWithName:@"LCTrie"
                                                          reason:@"key must not be empty"
                                                        userInfo:nil];
        @throw exception;
    }
}


-(LCTrieNode*)__nodeWithKey:(NSString*)key{
    [self __keyCheck:key];
    
    LCTrieNode * node = _root;
    NSUInteger len = key.length;
    for (int i = 0; i<len; i++) {
        if (node == nil || node.children == nil || node.children.count == 0) return nil;
        NSString * c = [key substringWithRange:NSMakeRange(i, 1)];
        node = [node.children objectForKey:c];
    }
    return node;
}

@end