jweinst1
7/19/2016 - 12:21 AM

trie implementation in swift

trie implementation in swift

//generic trie node
class TrieNode<T> {
    var matches:[String:T]
    var prefix:String
    var children:[String:TrieNode]
    var prefixLength:Int
    
    //static method that creates a trie such that all matches are linked to true, and have no linked value
    static func makeBoolTrie(strings:[String]) -> TrieNode<Bool>? {
        guard strings.count > 0 else {return nil}
        let bool_trie = TrieNode<Bool>(key: strings[0], value: true)
        for i in 1..<strings.count {
            bool_trie.insertKeyValue(strings[i], value: true)
        }
        return bool_trie
    }
    
    //static method that creates a trie of a generic linked type, needs a dictionary
    static func makeGenericTrie(zip:[(String, T)]) -> TrieNode<T>? {
        guard zip.count > 0 else {return nil}
        let gen_trie = TrieNode<T>(key: zip[0].0, value: zip[0].1)
        for (name, value) in zip.dropFirst() {
            gen_trie.insertKeyValue(name, value: value)
        }
        return gen_trie
    }
    
    
    
    //alternative initializer for recursive contrsuction
    //preflen defaults to zero for root node
    init(key:String, value:T, preflen:Int = 0) {
        self.matches = [String:T]()
        self.children = [String:TrieNode]()
        self.prefixLength = preflen
        self.prefix = String(Array(key.characters)[0..<self.prefixLength])
        self.matches[key] = value
        if key.characters.count != self.prefixLength {
            self[String(Array(key.characters)[0..<self.prefixLength+1])] = TrieNode(key: key, value: value, preflen: self.prefixLength+1)
        }
    }
    //recursive insertion into the trie, checks if child node already exists
    func insertKeyValue(key:String, value:T) {
        let child = self[String(Array(key.characters)[0..<self.prefixLength+1])]
        if String(Array(key.characters)[0..<self.prefixLength]) == self.prefix {
            self.addMatch(key, value: value)
            if child != nil {
                child!.insertKeyValue(key, value: value)
            }
        }
        else if child != nil {
            child!.insertKeyValue(key, value: value)
        }
        else {
            self[String(Array(key.characters)[0..<self.prefixLength+1])] = TrieNode(key: key, value: value, preflen: self.prefixLength+1)
        }
    }
    
    func getMatchesForString(searchkey:String) -> [String:T] {
        let empty = [String:T]()
        if searchkey == self.prefix {
            return self.matches
        }
        else {
            if let child = self[String(Array(searchkey.characters)[0..<self.prefixLength+1])] {
                return child.getMatchesForString(searchkey)
            }
            else {
                return empty
            }
        }
        
    }
    
    func addMatch(key:String, value:T) {
        self.matches[key] = value
    }
    //returns an array of the matches for the node
    func getMatchNames() -> [String] {
        return Array(self.matches.keys)
    }
    
    //checks if a node in the trie has a match
    func hasMatch(index:String) -> Bool {
        return self.matches[index] != nil
    }
    //checks if the current node is a leaf
    func isLeaf() -> Bool {
        return self.children.count == 0
    }
    
    subscript (index:String) -> TrieNode? {
        get {
            return children[index]
        }
        
        set(newValue) {
            children[index] = newValue
        }
    }
    
}

//example usage
/*var f = TrieNode<Bool>.makeBoolTrie(["fooo", "for", "foa"])
f.getMatchesForString("fo")*/
//["foa": true, "fooo": true, "for": true]

/*var testdict = [("foo", 5), ("fod", 3), ("foor", 1)]
var d = TrieNode.makeGenericTrie(testdict)
d!.getMatchesForString("foo")*/
//["foor": 1, "foo": 5]