Word Break Problem (Trie solution)

The problem, simply

You get a string s with no spaces, like catsanddog, and a dictionary of words. Can you put spaces into s so that every piece is a dictionary word? Words may be used many times. Say true or false.

Below, n is the length of s, K is the length of the longest dictionary word, and D is the total number of letters in the dictionary.

A story to remember it

Picture it

Examples

Way 1: the simple way

Try every first word. If the start of s is a dictionary word, cut it off and solve the rest the same way.

1. If nothing is left, the answer is true.
2. For every end position, if s[start..end) is in the dictionary and the rest can be split, the answer is true.
3. If no first word works, the answer is false.

import java.util.*;

class Solution {
    public static boolean wordBreak(String s, List<String> dict) {
        return canSplit(s, 0, new HashSet<>(dict));
    }

    private static boolean canSplit(String s, int start, Set<String> words) {
        if (start == s.length()) return true;               // nothing left to cut
        for (int end = start + 1; end <= s.length(); end++) {
            if (words.contains(s.substring(start, end))      // first piece is a word
                    && canSplit(s, end, words)) {            // and the rest works out
                return true;
            }
        }
        return false;
    }
}

How long it takes: the same rest of the string gets solved again and again. With many short words this tries close to every way of cutting, O(2^n). Extra room: the call stack, O(n).

Way 2: the clever way

Remember answers in a table, and use a trie to find all words that start at a spot in one walk.

1. Put every dictionary word in a trie.
2. Make dp of size n + 1. dp[k] is true when the first k letters can be cut into words. Set dp[0] = true.
3. For each i with dp[i] true, walk the trie with s[i], s[i+1], .... Each time you reach a word end at letter j, set dp[j + 1] = true. Stop when the trie has no branch for the next letter.
4. The answer is dp[n].

import java.util.*;

class Solution {
    private static class Node {
        Node[] next = new Node[26];
        boolean end;
    }

    public static boolean wordBreak(String s, List<String> dict) {
        Node root = new Node();
        for (String w : dict) {                          // build the trie
            Node cur = root;
            for (char ch : w.toCharArray()) {
                int k = ch - 'a';
                if (cur.next[k] == null) cur.next[k] = new Node();
                cur = cur.next[k];
            }
            cur.end = true;
        }

        int n = s.length();
        boolean[] dp = new boolean[n + 1];
        dp[0] = true;                                    // zero letters need zero words
        for (int i = 0; i < n; i++) {
            if (!dp[i]) continue;                        // no word may start here
            Node cur = root;
            for (int j = i; j < n; j++) {
                cur = cur.next[s.charAt(j) - 'a'];
                if (cur == null) break;                  // no word starts with s[i..j]
                if (cur.end) dp[j + 1] = true;           // s[i..j] is a word
            }
        }
        return dp[n];
    }
}

How long it takes: building the trie is O(D). From each of the n starts, the walk stops after at most K letters, because no word is longer. So O(n * K + D). Extra room: the dp table and the trie, O(n + 26 * D).

Watch out for

• Forgetting dp[0] = true. Without it nothing is ever marked and every answer is false.
• Stopping at the first word found from a start. cat and cats both start at 0, and both matter. Keep walking until the trie runs out.
• The answer lives in dp[n], not dp[n - 1]. The table is one longer than the string.

What interviewers ask next

• Print every way to cut the string? Keep, for each end, the list of starts that reach it, then build the sentences backward with backtracking. The output itself can be huge.
• Can you do it with a hash set instead of a trie? Yes. For each reachable i, check s.substring(i, j) for j up to i + K. Each check costs O(K) for hashing, so the trie is a bit faster and stops early.
• Memoized recursion instead of the table? Same idea from the other side: remember for each start whether the rest can be split. O(n * K) with the trie.
• Fewest words needed? Turn dp into numbers and keep the smallest count, or run a breadth first search over start positions.
• What if the dictionary is huge and s is short? Then the trie still costs only O(D) once, and each query is O(n * K).

Remember this

Grown-up words: the stones are positions in the string, the chalk marks are the 1D DP table, the word tree is a trie of the dictionary, and the jumps that spell words are valid segments.