LeetCode: 211. Design Add and Search Words Data Structure (and others)

// Solution @ Sergey Leschev

class TrieNode {
    var children: [Character: TrieNode]
    var isWord: Bool

    init() {
        children = [:]
        isWord = false
    }
}

class WordDictionary {
    var root: TrieNode

    init() {
        root = TrieNode()
    }

    func addWord(_ word: String) {
        var node = root
        for c in word {
            if let child = node.children[c] {
                node = child
            } else {
                let child = TrieNode()
                node.children[c] = child
                node = child
            }
        }
        node.isWord = true
    }

    func search(_ word: String) -> Bool {
        return searchHelper(Array(word), 0, root)
    }

    private func searchHelper(_ word: [Character], _ index: Int, _ node: TrieNode) -> Bool {
        if index == word.count {
            return node.isWord
        }

        let c = word[index]

        if c != "." {
            if let child = node.children[c] {
                return searchHelper(word, index+1, child)
            }
            return false
        } else {
            for child in node.children.values {
                if searchHelper(word, index+1, child) {
                    return true
                }
            }
            return false
        }
    }
}