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]