Exceeded time limit on Trie search with Backtracking using Swift

Question

I have been trying to solve Trie related problems over at leet code and came across this interesting problem 211. Design Add and Search Words Data Structure

Here is the question for convenience:

211. Design Add and Search Words Data Structure

Design a data structure that supports adding new words and finding if a string matches any previously added string.

Implement the WordDictionary class:

WordDictionary() Initializes the object.

void addWord(word) Adds word to the data structure, it can be matched later.

bool search(word) Returns true if there is any string in the data structure that matches word or false otherwise. word may contain dots '.' where dots can be matched with any letter.

Example:

Input

["WordDictionary","addWord","addWord","addWord","search","search","search","search"]

[[],["bad"],["dad"],["mad"],["pad"],["bad"],[".ad"],["b.."]] Output [null,null,null,null,false,true,true,true]

Explanation

WordDictionary wordDictionary = new WordDictionary();

wordDictionary.addWord("bad"); wordDictionary.addWord("dad");

wordDictionary.addWord("mad"); wordDictionary.search("pad"); // return

False wordDictionary.search("bad"); // return True

wordDictionary.search(".ad"); // return True

wordDictionary.search("b.."); // return True

Constraints:

1 <= word.length <= 25

word in addWord consists of lowercase English letters.

word in search consist of '.' or lowercase English letters.

There will be at most 3 dots in word for search queries.

At most 10⁴ calls will be made to addWord and search.

I got some help yesterday with a similar question on how to optimize my Trie implementation and tried my best to incorporate those into my solution to this question.

I got a solution to work with Trie and backtracking, however, I once again ran into the Time Limit exceeded conundrum.

I came across this Python solution on Youtube and tried to implement it using Swift, however that did not help.

Here is the search function:

private class WordDictionary {
    var root: TrieNode
    
    class TrieNode {
        var next = [Character: TrieNode]()
        var isWordEnd = false
    }
    
    init() {
        root = TrieNode()
    }
    
    func addWord(_ word: String) {
        var node = root
        
        for char in word {
            node = addInternal(char, at: node)
        }
        
        node.isWordEnd = true
    }
    
    func addInternal(_ char: Character, at node: TrieNode) -> TrieNode {
        if let nextNode = node.next[char] {
            return nextNode
        }
        
        let nextNode = TrieNode()
        node.next[char] = nextNode
        return nextNode
    }
    
    func search(_ word: String) -> Bool {
        
        func dfs(index: String.Index, at node: TrieNode) -> Bool {
            var currentNode = node
            var currentIndex = index
            
            while currentIndex < word.endIndex {
                let currentChar = word[currentIndex]
                
                if currentChar == "." {
                    if let nextNodes = Array(currentNode.next.values) as? [TrieNode] {
                        
                        for nextNode in nextNodes {
                            var nextIndex = currentIndex
                            word.formIndex(after: &nextIndex)
                            
                            if dfs(index: nextIndex, at: nextNode) {
                                return true
                            }
                        }
                    }
                    
                    return false
                }
                
                if let nextNode = currentNode.next[currentChar] {
                    currentNode = nextNode
                    word.formIndex(after: &currentIndex)
                    continue
                }
                
                return false
            }
            
            return currentNode.isWordEnd
        }
        
        return dfs(index: word.startIndex, at: root)
    }
}

This seems to be working correctly on these test cases:

private func test1() {
    let wordDictionary = WordDictionary()
    wordDictionary.addWord("baad")
    wordDictionary.addWord("dad")
    wordDictionary.addWord("mad")
    print(wordDictionary.search("baa")) // return False
    print(wordDictionary.search("bad")) // return False
    print(wordDictionary.search("baad")) // return True
    print(wordDictionary.search(".ad")) // return True
    print(wordDictionary.search("b..d")) // return True
    print(wordDictionary.search("pad")) // return False
    print(wordDictionary.search("bad")) // return False
    print(wordDictionary.search("b.ad")) // return True
    print(wordDictionary.search("mad")) // return True
    print(wordDictionary.search("bam")) // return False
}

private func test2() {
    let wordDictionary = WordDictionary()
    wordDictionary.addWord("a")
    wordDictionary.addWord("a")
    print(wordDictionary.search(".")) // return True
    print(wordDictionary.search("a")) // return True
    print(wordDictionary.search("aa")) // return False
    print(wordDictionary.search("a")) // return True
    print(wordDictionary.search(".a")) // return False
    print(wordDictionary.search("a.")) // return False
}

However, I get a time limit exceeded exception on this test case

I am not sure where the bottle neck and as at first glance, I cannot see any opportunity for memoization to improve the performance.

I am sure there might be other / more optimal solutions besides using a Trie, however, I am trying to get an optimal solution using a Trie since I am currently trying to improve my understanding of this topic.

Answer 1

Good job in incorporating the feedback from your previous question to improve the Trie implementation.

The dfs() function can be improved a bit: Creating an array of all successor nodes (in the case of a wildcard character)

if let nextNodes = Array(currentNode.next.values) as? [TrieNode] {
    for nextNode in nextNodes {
        // ...
    }
}

is not necessary, one can iterate over the values directly:

for nextNode in currentNode.next.values {
    // ...
}

The variable

var nextIndex = currentIndex
word.formIndex(after: &nextIndex)

can be made a constant with

let nextIndex = word.index(currentIndex, offsetBy: 1)

Also

if let nextNode = currentNode.next[currentChar] {
    currentNode = nextNode
    word.formIndex(after: &currentIndex)
    continue
}

return false

is in my opinion better written with an else block instead of continue:

if let nextNode = currentNode.next[currentChar] {
    currentNode = nextNode
    word.formIndex(after: &currentIndex)
} else {
    return false
}

Caching the results

In the test case that fails with a timeout error, the same string “ghkdnjmmgcddganjkigmabbho” is searched 5000 times. So caching does help. It can be easily added to the search function. Note that the cache must be invalidated if a new word is added to the dictionary:

class WordDictionary {
    var root: TrieNode
    var cache: [String: Bool] = [:]
    // ...    

    func addWord(_ word: String) {
        // ...
        cache.removeAll()
    }
    
    // ...    

    func search(_ word: String) -> Bool {
        if let result = cache[word] {
            return result
        }
        
        func dfs(index: String.Index, at node: TrieNode) -> Bool {
            // ...
        }
        
        let result = dfs(index: word.startIndex, at: root)
        cache[word] = result
        return result
    }
}

This can be optimized further:

• Repeated calls to cache.removeAll() can be avoided. Instead of emptying the cash in addWord(), set an “invalid” flag and empty the cache in the search() method if that flag is set.

• true entries in the cache remain valid if a new word is added to the dictionary. So instead of completely emptying the cache if a new word has been added, remove only the false entries.

The code then looks like this:

class WordDictionary {
    var root: TrieNode
    var cache: [String: Bool] = [:]
    var cacheValid = true
    
    class TrieNode {
        var next = [Character: TrieNode]()
        var isWordEnd = false
    }
    
    init() {
        root = TrieNode()
    }
    
    func addWord(_ word: String) {
        var node = root
        
        for char in word {
            node = addInternal(char, at: node)
        }
        
        node.isWordEnd = true
        cacheValid = false
    }
    
    func addInternal(_ char: Character, at node: TrieNode) -> TrieNode {
        if let nextNode = node.next[char] {
            return nextNode
        }
        
        let nextNode = TrieNode()
        node.next[char] = nextNode
        return nextNode
    }
    
    func search(_ word: String) -> Bool {
        if (!cacheValid) {
            let falseKeys = cache.compactMap { $0.value == false ? $0.key : nil }
            for key in falseKeys {
                cache[key] = nil
            }
            cacheValid = true
        } else if let result = cache[word] {
            return result
        }
        
        func dfs(index: String.Index, at node: TrieNode) -> Bool {
            var currentNode = node
            var currentIndex = index
            
            while currentIndex < word.endIndex {
                let currentChar = word[currentIndex]
                
                if currentChar == "." {
                    for nextNode in currentNode.next.values {
                        let nextIndex = word.index(currentIndex, offsetBy: 1)
                        if dfs(index: nextIndex, at: nextNode) {
                            return true
                        }
                    }
                    return false
                }
                
                if let nextNode = currentNode.next[currentChar] {
                    currentNode = nextNode
                    word.formIndex(after: &currentIndex)
                } else {
                    return false
                }
            }
            
            return currentNode.isWordEnd
        }
        
        let result = dfs(index: word.startIndex, at: root)
        cache[word] = result
        return result
    }
}