//
//  Trie.swift
//  Trie
//
//  Created by Neil Pankey on 3/6/15.
//  Copyright (c) 2015 Neil Pankey. All rights reserved.
//

public final class Trie<T> {
    private var value: T? = nil
    private var children: [Character:Trie<T>] = [:]

    public init() {
    }

    public func lookup(key: String) -> T? {
        if key.isEmpty {
            return value
        }

        let child = children[first(key)!]
        return child?.lookup(dropFirst(key)) ?? nil
    }

    public func insert(key: String, _ value: T) {
        if let c = first(key) {
            if var child = children[c] {
                child.insert(dropFirst(key), value)
            } else {
                var child = Trie()
                children[c] = child
                child.insert(dropFirst(key), value)
            }
        } else {
            self.value = value
        }
    }

    public func remove(key: String) {
        if let c = first(key) {
            if var child = children[c] {
                child.remove(dropFirst(key))
            }
            // TODO Error if the key doesn't exist?
        } else {
            value = nil
        }
    }
}

//
//  Trie.swift
//  Trie
//
//  Created by Neil Pankey on 3/6/15.
//  Copyright (c) 2015 Neil Pankey. All rights reserved.
//

public final class Trie<T> {
    private var value: T? = nil
    private var children: [Character:Trie<T>] = [:]

    public init() {
    }

    public func lookup(key: String) -> T? {
        if key.isEmpty {
            return value
        }

        let child = children[first(key)!]
        let subkey = key[key.startIndex.successor()..<key.endIndex]
        return child?.lookup(subkey) ?? nil
    }

    public func insert(key: String, _ value: T) {
        if let c = first(key) {
            let subkey = key[key.startIndex.successor()..<key.endIndex]
            if var child = children[c] {
                child.insert(subkey, value)
            } else {
                var child = Trie()
                children[c] = child
                child.insert(subkey, value)
            }
        } else {
            self.value = value
        }
    }

    public func remove(key: String) {
        if let c = first(key) {
            if var child = children[c] {
                child.remove(key[key.startIndex.successor()..<key.endIndex])
            }
            // TODO Error if the key doesn't exist?
        } else {
            value = nil
        }
    }
}

--- 
+++ 
@@ -19,19 +19,17 @@
         }
 
         let child = children[first(key)!]
-        let subkey = key[key.startIndex.successor()..<key.endIndex]
-        return child?.lookup(subkey) ?? nil
+        return child?.lookup(dropFirst(key)) ?? nil
     }
 
     public func insert(key: String, _ value: T) {
         if let c = first(key) {
-            let subkey = key[key.startIndex.successor()..<key.endIndex]
             if var child = children[c] {
-                child.insert(subkey, value)
+                child.insert(dropFirst(key), value)
             } else {
                 var child = Trie()
                 children[c] = child
-                child.insert(subkey, value)
+                child.insert(dropFirst(key), value)
             }
         } else {
             self.value = value
@@ -41,7 +39,7 @@
     public func remove(key: String) {
         if let c = first(key) {
             if var child = children[c] {
-                child.remove(key[key.startIndex.successor()..<key.endIndex])
+                child.remove(dropFirst(key))
             }
             // TODO Error if the key doesn't exist?
         } else {
Use dropFirst instead of slicing
