Trie autocomplete

TrieAutocomplete.swift

import Foundation

/*

    D
     O
      O
    G   R
          S
    
    look for prefix, then DFS children for option

*/

import Foundation

/*

    D
     O
      O
    G   R
          S
    
    look for prefix, then DFS children for option

*/

class Trie {
    class Node {
        var children: [Character: Node]
        var isWord: Bool
        
        init(isWord: Bool = false) {
            children = [:]
            self.isWord = isWord
        }
    }
    
    private let root: Node
    
    // build a trie
    init(_ words: [String]) {        
        root = Node()
        for word in words {
            var current = root
            for character in word {
                if let child = current.children[character] {
                    // existing node
                    current = child
                } else {
                    // new node
                    let n = Node()
                    current.children[character] = n
                    current = n
                }
            }
            // that's the whole word node
            current.isWord = true
        }
    }

    func autocomplete(_ prefix: String) -> [String] {
        var current = root
        // traverse over the prefix word
        for character in prefix {
            if let child  = current.children[character] {
                current = child
            } else {
                return []
            }
        }
        // then add "combinations" by traversing with dfs
        return dfs(current, prefix)
    }
    
    private func dfs(_ node: Node, _ prefix: String) -> [String] {
        var words: [String] = []
        if node.isWord {
            words += [prefix]
        }
        
        for (character, child) in node.children {
            words += dfs(child, prefix + String(character))
        }
        return words
    }
}


let t = Trie(["dog","dark","cat","door","dodge"])
let list = t.autocomplete("do")
print(list)