Determine all ways a string can be split into valid words, given a dictionary of all words

Question

This function word_break(s, dictionary) receives a string s and a set of all valid words dictionary, and must return all possible ways the string s can be split into valid words, if any.

For example, a string catsanddog with dictionary {'cat', 'cats', 'sand', 'and', 'dog'} can be broken in two ways: cats and dog and cat sand dog.

My solution to this recursively breaks the input string at various index (offset), and uses memoization as I realized we may be breaking at the same offset multiple times.

Here is the full code along with a test case, in Python 2.x:

def memoize(fn):
    memo = {}
    def wrapped_fn(*args):
        if args not in memo:
            memo[args] = fn(*args)
        return memo[args]
    return wrapped_fn

def word_break(s, dictionary):    
    @memoize
    def _word_break(offset):
        if offset == len(s):
            return []

        breaks_from_offset = []
        if s[offset:] in dictionary:
            breaks_from_offset.append([s[offset:]])

        for i in range(offset+1, len(s)):
            if s[offset:i] in dictionary:
                for break_from_i in _word_break(i):
                    breaks_from_offset.append([s[offset:i]] + break_from_i)

        return breaks_from_offset

    return [' '.join(words) for words in _word_break(0)]

# print word_break('catsanddog', {'cat', 'cats', 'sand', 'and', 'dog'})
# -> [cat sand dog', 'cats and dog']

Based on LeetCode question Word Break II: https://leetcode.com/problems/word-break-ii/description/

As this is an "interview" algorithm question, I would love feedback that keeps that setting in mind. That said, I also appreciate any feedback to help me write more production-ready Python code :)

Answer 1

Honestly it looks like you're doing this wrong. Currently you're splitting the word into \$n^2\$ words and checking if they are in the dictionary. Rather than doing this you can traverse a Trie and loop through the input once. To do this you'd just need a 'partial words list' which is a set of Trie nodes. You then just need to make some lists for the output, and then return the one that you think is best. This can be a bit tricky to do, but the challenge is finding which output list(s) to append to, which you can do when adding new tries to the partial words list.

I'd implement this in the following way:

import collections.abc


class TrieNode(collections.abc.MutableMapping):
    def __init__(self, k):
        self._data = {}
        self._value = k
        self.end = False

    @staticmethod
    def build(iterable):
        root = TrieNode(None)
        for key in iterable:
            root[key] = True
        return root

    @property
    def value(self):
        return self._value

    def _keys(self, key):
        partial = ''
        for k in key:
            partial += k
            yield k, partial

    def _walk(self, data, key, *, build=False):
        if not key:
            raise ValueError()

        node = data
        if not build:
            for k in key[:-1]:
                node = node._data[k]
        else:
            for k, key_ in self._keys(key[:-1]):
                node = node._data.setdefault(k, TrieNode(key_))
        return key[-1], node

    def __getitem__(self, key):
        key, node = self._walk(self, key)
        return node._data[key]

    def __setitem__(self, key, value):
        k, node = self._walk(self, key, build=True)
        node = node._data.setdefault(k, TrieNode(key))
        node.end = value

    def __delitem__(self, key):
        key, node = self._walk(self, key)
        del node._data[key]

    def __iter__(self):
        return iter(self._data)

    def __len__(self):
        return len(self._data)


def word_break(word_string, words):
    words = TrieNode.build(words)
    output = {0: [[]]}
    partials = []
    for i, k in enumerate(word_string, 1):
        new_partials = []
        for partial in partials + [words]:
            partial = partial.get(k)
            if partial is None:
                continue

            new_partials.append(partial)
            if not partial.end:
                continue

            val = partial.value
            prevs = output.get(i - len(val))
            if prevs is not None:
                output.setdefault(i, []).extend([p + [val] for p in prevs])
        partials = new_partials
    return output[len(word_string)]

if __name__ == '__main__':
    words = word_break('catsanddog', {'cat', 'cats', 'sand', 'and', 'dog'})
    print(words)