Return to Answer

• 24.2k
• 2
• 38
• 96

• 24.2k
• 2
• 38
• 96

Summary: Your code is written clearly, and works correctly as far as I can see. It can be improved at a few points, and string slices can be used for a small performance improvement. In order to make it work with really large strings within the given time constraint, additional techniques like memoization must be used.

Simplify and improve the code

The var data: Character? of class TrieNode is assigned, but never used. It can be removed.

In func insertInternal() a

var nextNode = TrieNode()

is created but not used if node.next[char] already exists. Also the forced unwrapping can be avoided:

func insertInternal(_ char: Character, at node: TrieNode) -> TrieNode {
 if let nextNode = node.next[char] {
 return nextNode
 } else {
 let nextNode = TrieNode()
 node.next[char] = nextNode
 return nextNode
 }
}

In func doesWordExist the variable name of

let nextNode = rootNode.next

is misleading: it is not a node but a dictionary. Again the forced unwrapping (and the comparison against nil) can be avoided:

func doesWordExist(_ word: Substring) -> Bool {
 var rootNode = root
 for char in word {
 if let nextNode = rootNode.next[char] {
 rootNode = nextNode
 } else {
 return false
 }
 }
 return rootNode.isWordEnd
}

Use Substring

It is more efficient to pass Substrings to the recursively called wordBreakInternal() functions instead of Strings. Substring is a "string slice" and shares the storage with the original string. The helper function is then declared as

func wordBreakInternal(_ s: Substring, wordDict: Trie) -> Bool

and called as

return wordBreakInternal(s[...], wordDict: trie)

This makes the function a bit faster, but not fast enough to solve the programming challenge within the time limit.

Memoization

The function is too slow for long strings because of the heavy recursion. One approach to reduce the amount of recursive calls is memoization: For each substring which has been determined to be word-breakable or not, remember the result:

func wordBreak(_ s: String, _ wordDict: [String]) -> Bool {
 let trie = Trie()
 var memo = [Substring: Bool]()
 
 for word in wordDict {
 trie.insert(word)
 }
 
 func wordBreakInternal(_ s: Substring) -> Bool {
 if let memoized = memo[s] {
 return memoized
 }
 
 if s.isEmpty { return true }
 
 let leftIndex = s.startIndex
 var rightIndex = s.startIndex
 
 while rightIndex < s.endIndex {
 let word = (s[leftIndex ... rightIndex])
 
 if trie.doesWordExist(word) {
 var nextStartIndex = rightIndex
 s.formIndex(after: &nextStartIndex)
 let nextWord = (s[nextStartIndex ..< s.endIndex])
 
 if wordBreakInternal(nextWord) {
 memo[s] = true
 return true
 }
 }
 s.formIndex(after: &rightIndex)
 }
 
 memo[s] = false
 return false
 }
 
 return wordBreakInternal(s[...])
}

I have also made the helper function a nested function in the main function, so that the memoization dictionary (and the trie) need not be passed around as function parameters.

This function computes

let s = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaab"
let wordDict = ["a","aa","aaa","aaaa","aaaaa","aaaaaa","aaaaaaa","aaaaaaaa","aaaaaaaaa","aaaaaaaaaa"]
let result = wordBreak(s, wordDict)

in about 7 milliseconds (compiled in Release mode, running on a Macbook Air M2).

To trie or not to trie

Use have used a Trie in order to determine quickly if the strings starts with one of the given words. As it turns out, using the built-in hasPrefix() method does this even better. Using string slices and memoization as before we get

func wordBreak(_ s: String, _ wordDict: [String]) -> Bool {
 var memo = [Substring: Bool]()
 
 func wordBreakInternal(_ s: Substring) -> Bool {
 if let memoized = memo[s] {
 return memoized
 }
 if s.isEmpty {
 return true
 }
 for word in wordDict {
 if s.hasPrefix(word) {
 if wordBreakInternal(s.dropFirst(word.count)) {
 memo[s] = true
 return true
 }
 }
 }
 memo[s] = false
 return false
 
 }
 return wordBreakInternal(s[...])
}

which computes the above test case in less than a millisecond.