Problem Statement in English

You’re given a list of words (each word being a string) and you need to implement a WordFilter class that allows for searching through the words using a given prefix and suffix.

Approach

We can solve this problem using a Trie. We store every word in a prefix tree, and its reverse in a suffix tree.

Also, each Trie node has an extra references parameter, which is a Set, in which we store the all the words that sequence belongs to.

Solution in Python


# define the trie node
class TrieNode:
    def __init__(self) -> None:
        self.children: dict[str, TrieNode] = {}
        self.references: set[str] = set()

    # adds a word to the trie
    def addWord(self, word):
        ptr = self
        for c in word:
            if c not in ptr.children:
                ptr.children[c] = TrieNode()
            temp = ptr.children[c]
            temp.references.add(word)
            ptr = temp

    # adds a word to the trie in reverse
    def addWordReverse(self, word):
        ptr = self
        for c in reversed(word):
            if c not in ptr.children:
                ptr.children[c] = TrieNode()
            temp = ptr.children[c]
            temp.references.add(word)
            ptr = temp


class WordFilter:

    def __init__(self, words: List[str]):
        # define the prefix trie
        self.prefixTree = TrieNode()

        # define the suffix trie
        self.suffixTree = TrieNode()

        # define the index
        #
        # we use this to store
        # the index of the word in `words`
        self.index: dict[str, int] = {}

        # add the words to the trie
        for i in range(len(words)):
            # add the word to the prefix trie
            self.prefixTree.addWord(words[i])

            # add the word in reverse to the suffix trie
            self.suffixTree.addWordReverse(words[i])

            # add the word to the index
            self.index[words[i]] = i

    def f(self, pref: str, suff: str) -> int:

        # search for the word in the prefix trie
        ptr = self.prefixTree
        for c in pref:
            if c in ptr.children:
                ptr = ptr.children[c]
            else:
                return -1
        # save the set of all indices where this prefix exists
        prefixRes = ptr.references

        # search for the word in the suffix trie
        ptr = self.suffixTree
        # iterate over the word in reverse
        # so that we can add it in reverse
        for c in reversed(suff):
            if c in ptr.children:
                ptr = ptr.children[c]
            else:
                return -1
        # save the set of all indices where this suffix exists
        suffixRes = ptr.references

        found = -1
        # check to look for the maximum index that is common in both sets
        for ref in prefixRes:
            if ref in suffixRes:
                # store the maximum index
                found = max(found, self.index[ref])

        # return the maximum index, or -1 if no common index exists
        return found

Complexity

  • Time: O(n)O(n)
    Where nn is the number of characters in prefix and suffix.

  • Space: O(n)O(n)
    Since we store the words in the trie, where nn is the number of all characters in the words list.

Mistakes I Made

I actually tried using a SortedSet to get a constant time lookup of the maximum index, but that resulted in poor insertion time (lognlog n per insertion), and also took up a bunch of space. So I decided to use a simple set instead.

And we are done.