PrefixTrie Cleanup

TrieCleanup.swift

public struct PrefixTrie<Key: Collection, Value> where Key.Element: Hashable {
    internal class Node {
        var value: Value?
        var step: Key.Element?
        var next = [Key.Element: Node]()
        var parent: Node?

        init(value: Value?, step: Key.Element? = nil, parent: Node? = nil) {
            self.value = value
            self.parent = parent
            self.step = step
        }
    }

    /// Creates a prefix trie with nothing inside.
    public init() {}

    // Start with an empty node at the root, to make traversal easy
    var root = Node(value: nil)

    /// Finds the value associated with the longest prefix that matches the
    /// provided key.
    /// For example, for a trie that has `["Help", "Hello", "Helping"]` inside,
    /// Searching for `"Helper"` gives you the value associated with `"Help"`,
    /// while searching for `"Helping a friend"` gives you `"Helping"`.
    public subscript(key: Key) -> Value? {
        get {
            var bestMatch: Value?
            var current = root
            for elt in key {
                guard let next = current.next[elt] else { break }
                current = next
                if let value = current.value {
                    bestMatch = value
                }
            }
            return bestMatch
        }
        set {
            var index = key.startIndex
            var current = root

            // Traverse as much of the prefix as we can, keeping track of the index
            // we ended on
            while index < key.endIndex, let next = current.next[key[index]] {
                key.formIndex(after: &index)
                current = next
            }

            // We're matching a prefix of an existing key in the trie
            if index == key.endIndex {
                // Update the value in the trie with the new value

                current.value = newValue

                var currentNode: Node! = current
                var parentNode = current.parent

                while parentNode != nil && currentNode.step != nil {

                    if parentNode?.next[currentNode.step!]?.value == nil && parentNode?.next[currentNode.step!]?.next.keys.count == 0 {

                        parentNode?.next[currentNode.step!] = nil
                        currentNode = parentNode
                        parentNode = parentNode?.parent
                    }
                    else {
                        parentNode = nil
                    }
                }

                return
            }

            while index < key.endIndex {
                // Fill out the remaining nodes with nils, until we get to the last
                // element
                let step = key[index]

                key.formIndex(after: &index)
                let new = Node(value: index == key.endIndex ? newValue : nil, step: step, parent: current)
                current.next[step] = new
                current = new
            }
        }
    }

}


func count(_ trie: PrefixTrie<String, String>) -> Int {

    var n: PrefixTrie.Node? = trie.root

    var count = 0

    var nodes = [PrefixTrie<String, String>.Node]()
    var keys = [Character]()

    while n != nil {

        nodes.append(contentsOf: n!.next.values)
        keys.append(contentsOf: n!.next.keys)
        print(keys)
        count += n!.next.keys.count

        n = nodes.popLast()
    }

    return count
}


var trie = PrefixTrie<String, String>()
trie["hello"] = "hello"
trie["help"] = "help"
trie["halo"] = "halo"
trie["helping"] = "helping"


count(trie)

trie["helping"] = nil
trie["halo"] = nil

count(trie)