Prefix and Suffix Search

How it works

A Trie is ideal for prefix matching, but we also need suffix matching. The trick is to insert modified strings that encode both. For each word, we insert all rotations of the form suffix + "{" + prefix into the Trie. The character { is chosen because it comes right after z in ASCII, giving us a unique separator. When querying, we search for suff + "{" + pref, which finds words matching both conditions simultaneously.

Steps

1. Build a Trie where each node stores children (26 letters + 1 separator) and the latest word index.
2. For each word at index i:
   • For every suffix starting position j:
     • For every prefix ending position k:
       • Insert suffix + "{" + prefix into the Trie.
       • Mark each node along the path with index i.
3. For f(pref, suff):
   • Search the Trie for suff + "{" + pref.
   • Return the index stored at the final node, or -1 if not found.

Complexity

• Time complexity:
  • O(n * m ^ 3) time for initialization.
  • O(m) for each f() function call.
• Space complexity: O(n * m ^ 3)

Where n is the number of words and m is the average length of each word.

Things to watch for

Choosing the Wrong Separator Character

When combining prefix and suffix into a single key (e.g., suffix + "{" + prefix), you must use a character that cannot appear in the input. Since the problem specifies only lowercase letters a-z, using { (ASCII 123, right after z) is safe. Using a character like _ or any letter would cause collisions with valid input.

Returning the First Match Instead of the Highest Index

The problem asks for the word with the highest index that matches both prefix and suffix. If multiple words satisfy the condition, you must return the largest index. When iterating forward and storing in a hash map, later words naturally overwrite earlier ones. When iterating in brute force, you must iterate from the end backward.

Off-by-One Errors in Prefix/Suffix Generation

When generating all prefixes, ensure you include both the empty prefix and the full word. Similarly for suffixes. Missing the empty string or the complete word means some valid queries will fail. The loop bounds should be 0 to word.length() inclusive for prefixes and suffixes.

Incorrect Trie Node Index Updates

In the Trie approach, each node should store the latest (highest) word index seen along that path. A common mistake is only updating the index at the leaf node. You must update the index at every node along the insertion path so that partial matches correctly return the highest index.

Memory Limit Exceeded from Excessive Key Generation

The hash map and Trie approaches generate O(m^2) entries per word (all prefix-suffix combinations), leading to O(n × m^3) total space when accounting for string storage. For large inputs, this can exceed memory limits. Ensure your implementation reuses string builders or avoids unnecessary string copies during key construction.

from typing import List

class TrieNode:
    def __init__(self):
        self.children = [None] * 27
        self.index = -1

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def addWord(self, w, i):
        cur = self.root
        for ch in w:
            c = ord(ch) - ord('a')
            if not cur.children[c]:
                cur.children[c] = TrieNode()
            cur = cur.children[c]
        cur.index = i

    def search(self, w):
        cur = self.root
        for ch in w:
            c = ord(ch) - ord('a')
            if not cur.children[c]:
                return -1
            cur = cur.children[c]
        return cur.index

class WordFilter:
    def __init__(self, words: List[str]):
        self.trie = Trie()
        self.CHAR = '{'
        for i, w in enumerate(words):
            w_len = len(w)
            for j in range(w_len):
                suffix = w[j:]
                for k in range(w_len + 1):
                    prefix = w[:k]
                    self.trie.addWord(suffix + self.CHAR + prefix, i)

    def f(self, pref: str, suff: str) -> int:
        return self.trie.search(suff + self.CHAR + pref)

class TrieNode {
public:
    TrieNode* children[27];
    int index;

    TrieNode() {
        for (int i = 0; i < 27; i++) {
            children[i] = nullptr;
        }
        index = -1;
    }
};

class Trie {
private:
    TrieNode* root;

public:
    Trie() {
        root = new TrieNode();
    }

    void addWord(const string& word, int i) {
        TrieNode* cur = root;
        for (char ch : word) {
            int c = (ch == '{') ? 26 : (ch - 'a');
            if (cur->children[c] == nullptr) {
                cur->children[c] = new TrieNode();
            }
            cur = cur->children[c];
        }
        cur->index = i;
    }

    int search(const string& word) {
        TrieNode* cur = root;
        for (char ch : word) {
            int c = (ch == '{') ? 26 : (ch - 'a');
            if (cur->children[c] == nullptr) {
                return -1;
            }
            cur = cur->children[c];
        }
        return cur->index;
    }
};

class WordFilter {
private:
    Trie trie;
    const char CHAR = '{';

public:
    WordFilter(vector<string>& words) {
        for (int i = 0; i < words.size(); i++) {
            string word = words[i];
            int wLen = word.length();
            for (int j = 0; j < wLen; j++) {
                string suffix = word.substr(j);
                for (int k = 0; k <= wLen; k++) {
                    string prefix = word.substr(0, k);
                    trie.addWord(suffix + CHAR + prefix, i);
                }
            }
        }
    }

    int f(string pref, string suff) {
        return trie.search(suff + CHAR + pref);
    }
};