Word Break

The problem description can be found in LeetCode.

Approach

A brute force approach for this problem would consist of consuming the string s by starting from any word of wordDict. Of course, this would require common characters between the strings. In fact, we would not complete any search with an input like:

s = "hello"
wordDict = ["foo", "bar"]

This would suggest looking any character for the strings in wordDict and start to match them with the characters in s. If we get a complete match for a certain string w in wordDict, we can start to look at the wordDict again with s[:len(w)].

This is a time consuming approach because it would require to go through all the characters in s and then all the characters of each string of wordDict.

However, we should already know about a data structure that can quickly look at prefixes: the Trie. So, we could store all the words in wordDict in a Trie and just iterate over the strings s to look for common characters.

Solution

from typing import Dict, List, Optional


class TrieNode:
    def __init__(self, is_word: Optional[bool] = False) -> None:
        self.children: Dict[str, "TrieNode"] = {}
        self.is_word = is_word


class Trie:
    def __init__(self) -> None:
        self.head = TrieNode()

    def insert(self, word: str) -> None:
        p = self.head

        for char in word:
            if char not in p.children:
                p.children[char] = TrieNode()

            p = p.children[char]

        p.is_word = True

    def search(self, word: str) -> bool:
        return self._contains(word=word, match_word=True)

    def startsWith(self, prefix: str) -> bool:
        return self._contains(word=prefix, match_word=False)

    def _contains(self, word: str, match_word: bool) -> bool:
        p = self.head

        for char in word:
            if char not in p.children:
                return False

            p = p.children[char]

        return p.is_word if match_word else True


class Solution:
    def wordBreak(self, s: str, wordDict: List[str]) -> bool:
        def dfs(i: int, trie: Trie, memo: Dict[int, bool]) -> bool:
            if i in memo:
                return memo[i]
            if i == len(s):
                return True
            if i > len(s):
                return False

            p = trie.head
            for j in range(i, len(s)):
                char = s[j]
                if char not in p.children:
                    memo[i] = False
                    return memo[i]

                p = p.children[char]
                if p.is_word and dfs(j + 1, trie, memo):
                    memo[i] = True
                    return memo[i]

            memo[i] = False
            return memo[i]

        trie = Trie()
        for word in wordDict:
            trie.insert(word)

        return dfs(0, trie, {})

Complexity

• Time: $O(mw+n^2)$
• Space: $O(mw+n)$

Where:

• $m$ is the length of the list wordList.
• $w$ is the length of the longest string in wordList.
• $n$ is the length of the string s.