Simple Trie.swift

class TrieNode {
  var key: Character?
  var children: [Character: TrieNode] = [:]
  weak var parent: TrieNode?
  var end: Bool
  init(_ key: Character?) {
    self.key = key
    self.children = [:]
    self.parent = nil
    self.end = false
  }
  
  func getWords() -> [Character] {
    var output:[Character] = []
    var node:TrieNode? = self
    while node != nil {
      if let c = node?.key {
        output.append(c)
      }
      node = node!.parent
    }
    
    return output
  }
}

class Trie {
  let root = TrieNode(nil)
  
  func insert(_ word: [Character]) {
    if word.count == 0 {
      return
    }
    var node: TrieNode = self.root
    var c: Character
    for i in 0 ..< word.count {
      c = word[i]
      if node.children[c] == nil {
        let child = TrieNode(c)
        child.parent = node
        node.children[c] = child
        node = child
      } else {
        node = node.children[c]!
      }
      
      if i == word.count - 1 {
        node.end = true
      }
    }
  }
  
  func contains(word: [Character]) -> Bool {
    var node: TrieNode = self.root
    for i in 0 ..< word.count {
      let c = word[i]
      guard let childNode = node.children[c] else {
        return false
      }
      node = childNode
    }
    return node.end
  }
  
  class func findAllWords(node: TrieNode, output: inout [[Character]]) {
    if node.end {
      output.insert(node.getWords().reversed(), at: 0)
    }
    for child in node.children {
      findAllWords(node: child.value, output: &output)
    }
  }
  
  func find(word: [Character]) -> [[Character]] {
    var output:[[Character]] = []
    var node: TrieNode = self.root
    
    // seek the last node of word
    for i in 0 ..< word.count {
      let c = word[i]
      guard let childNode = node.children[c] else {
        return output // trie does not contains whole word
      }
      node = childNode
    }
    Self.findAllWords(node: node, output: &output)
    
    return output
  }
}

let aTrie = Trie()
let word1 = [Character]("hello")
let word2 = [Character]("hellium")
let word3 = [Character]("ell")
aTrie.insert(word1)
aTrie.insert(word2)
aTrie.insert(word3)

aTrie.contains(word: word1) // true
aTrie.contains(word: word2)
aTrie.contains(word: [Character]("hell")) // false
aTrie.contains(word: [Character]("jellium")) // false

//let keyword = [Character]("hello")
let keyword = [Character]("e")
aTrie.find(word: keyword) // [["h", "e", "l", "l", "o"]]