Prefix and Suffix Search

---

Prefix and Suffix Search: Design a special dictionary that searches the words in it by a prefix and a suffix.

Implement the WordFilter class:

• WordFilter(string[] words) Initializes the object with the words in the dictionary.
• f(string pref, string suff) Returns the index of the word in the dictionary, which has the prefix pref and the suffix suff. If there is more than one valid index, return the largest of them. If there is no such word in the dictionary, return -1.

Example 1:

Input
["WordFilter", "f"]
[[["apple"]], ["a", "e"]]
Output
[null, 0]
Explanation
WordFilter wordFilter = new WordFilter(["apple"]);
wordFilter.f("a", "e"); // return 0,
because the word at index 0 has prefix = "a" and suffix = "e".

---

Constraints:

• 1 <= words.length <= 10^4
• 1 <= words[i].length <= 7
• 1 <= pref.length, suff.length <= 7
• words[i], pref and suff consist of lowercase English letters only.
• At most 10^4 calls will be made to the function f.

---

Try this Problem on your own or check similar problems:

1. Design Add and Search Words Data Structure

Solution:

Java

class WordFilter {
  class Trie {
      public class TrieNode{
        TrieNode[] children;
        int weight;
        public TrieNode(){
            children = new TrieNode[27];
            weight = 0;
        }
      }

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

      public void insert(String word, int weight) {
        for (int i = 0; i < word.length(); ++i) {
          TrieNode iter = root;
          iter.weight = weight;
          for (int j = i; j < 2 * word.length() - 1; ++j) {
            int k = word.charAt(j % word.length()) - 'a';
            if (iter.children[k] == null)
              iter.children[k] = new TrieNode();
            iter = iter.children[k];
            iter.weight = weight;
          }
        }
      }

      public int search(String word){
          TrieNode iter = root;
          for (int i = 0; i < word.length(); ++i) {
            char letter = word.charAt(i);
            if (iter.children[letter - 'a'] == null) return -1;
            iter = iter.children[letter - 'a'];
          }
          return iter.weight;
      }

  }
  Trie trie = new Trie();

  public WordFilter(String[] words) {
      for (int i = 0; i < words.length; ++i) {
        trie.insert(words[i] + '{', i);
      }
  }

  public int f(String pref, String suff) {
      return trie.search(suff + '{' + pref);
  }
}

/**
 * Your WordFilter object will be instantiated and called as such:
 * WordFilter obj = new WordFilter(words);
 * int param_1 = obj.f(pref,suff);
 */

JavaScript

/**
 * @param {string[]} words
 */
var WordFilter = function (words) {
  this.trie = new Trie();

  for (let i = 0; i < words.length; i++) {
    this.trie.insert(words[i] + "{", i);
  }
};

/**
 * @param {string} pref
 * @param {string} suff
 * @return {number}
 */
WordFilter.prototype.f = function (pref, suff) {
  return this.trie.search(suff + "{" + pref);
};

class Trie {
  constructor() {
    this.root = new TrieNode();
  }

  insert(word, weight) {
    for (let i = 0; i < word.length; i++) {
      let iter = this.root;
      iter.weight = weight;
      for (let j = i; j < 2 * word.length - 1; j++) {
        const k =
          word.charAt(j % word.length) === "{"
            ? 26
            : word.charAt(j % word.length).charCodeAt(0) - 97;
        if (iter.children[k] === null) {
          iter.children[k] = new TrieNode();
        }
        iter = iter.children[k];
        iter.weight = weight;
      }
    }
  }

  search(word) {
    let iter = this.root;
    for (let i = 0; i < word.length; i++) {
      const letter = word.charAt(i);
      const k = letter === "{" ? 26 : letter.charCodeAt(0) - 97;
      if (iter.children[k] === null) return -1;
      iter = iter.children[k];
    }
    return iter.weight;
  }
}

class TrieNode {
  constructor() {
    this.children = Array(27).fill(null);
    this.weight = 0;
  }
}

/**
 * Your WordFilter object will be instantiated and called as such:
 * var obj = new WordFilter(words)
 * var param_1 = obj.f(pref,suff)
 */

Python

class WordFilter:

    def __init__(self, words: List[str]):
      self.trie = Trie()

      for i in range(len(words)):
        self.trie.insert(words[i] + '{', i)


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

class Trie:
  class TrieNode:
    def __init__(self):
      self.children = [None] * 27
      self.weight = 0

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

  def insert(self, word, weight):
    for i in range(len(word)):
      iter = self.root
      iter.weight = weight
      for j in range(i, 2 * len(word) - 1):
        k = ord(word[j % len(word)]) - ord('a')
        if iter.children[k] is None:
          iter.children[k] = self.TrieNode()
        iter = iter.children[k]
        iter.weight = weight

  def search(self, word):
    iter = self.root
    for i in range(len(word)):
      letter = word[i]
      if iter.children[ord(letter) - ord('a')] is None:
        return -1
      iter = iter.children[ord(letter) - ord('a')]
    return iter.weight



# Your WordFilter object will be instantiated and called as such:
# obj = WordFilter(words)
# param_1 = obj.f(pref,suff)

C++

class Trie {
private:
  struct TrieNode {
    TrieNode* children[27];
    int weight;

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

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

  void insert(string word, int weight) {
    for (int i = 0; i < word.length(); i++) {
      TrieNode* iter = root;
      iter->weight = weight;
      for (int j = i; j < 2 * word.length() - 1; j++) {
        int k = word[j % word.length()] - 'a';
        if (iter->children[k] == nullptr) {
          iter->children[k] = new TrieNode();
        }
        iter = iter->children[k];
        iter->weight = weight;
      }
    }
  }

  int search(string word) {
    TrieNode* iter = root;
    for (int i = 0; i < word.length(); i++) {
      char letter = word[i];
      if (iter->children[letter - 'a'] == nullptr) return -1;
      iter = iter->children[letter - 'a'];
    }
    return iter->weight;
  }
};

class WordFilter {
private:
  Trie trie;
public:
    WordFilter(vector<string>& words) {
      for (int i = 0; i < words.size(); i++) {
        trie.insert(words[i] + '{', i);
      }
    }

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

/**
 * Your WordFilter object will be instantiated and called as such:
 * WordFilter* obj = new WordFilter(words);
 * int param_1 = obj->f(pref,suff);
 */

---

Time/Space Complexity:

• Time Complexity: O(nk^2 + Qk)
• Space Complexity: O(nk^2)

---

Explanation:

We can have two tries one for inserting all prefixes and the other inserting suffixes. We also track indices of words Set<Integer> wordIdx for each of the nodes, so to find a word with maximum index with certain prefix and suffix we can search for prefix and suffix in respective tries and find the intersection of the two sets of words (which start with prefix, and end in suffix) and return maximum element of that intersection. This will lead to TLE if the sets become large.

Java

class WordFilter {
    class Trie {
      public class TrieNode{
          TrieNode[] children;
          boolean isWord = false;
          Set<Integer> wordIdx;
          public TrieNode(){
              children = new TrieNode[26];
              wordIdx = new HashSet<>();
          }
      }

      TrieNode suffixRoot, prefixRoot;
      public Trie() {
          suffixRoot = new TrieNode();
          prefixRoot = new TrieNode();
      }

      public void insert(String word, int idx, boolean isSuffixRoot) {
          TrieNode iter = isSuffixRoot ? suffixRoot : prefixRoot;
          for(int i = 0; i < word.length(); ++i){
              char letter = word.charAt(i);
              if(iter.children[letter - 'a'] == null){
                  iter.children[letter - 'a'] = new TrieNode();
              }
              iter = iter.children[letter - 'a'];
              iter.wordIdx.add(idx);
          }
          iter.isWord = true;
          iter.wordIdx.add(idx);
      }

      private TrieNode searchPrefix(String prefix, boolean isSuffixRoot){
          TrieNode iter = isSuffixRoot ? suffixRoot : prefixRoot;
          for(int i = 0; i < prefix.length(); ++i){
              char letter = prefix.charAt(i);
              if(iter.children[letter - 'a'] == null){
                  return null;
              }
              iter = iter.children[letter - 'a'];
          }
          return iter;
      }

      public Set<Integer> startsWith(String prefix, boolean suffixLookup) {
          TrieNode node = searchPrefix(prefix, suffixLookup);
          return node == null ? null : node.wordIdx;
      }
    }

    Trie trie = new Trie();

    private String reverseString(String s){
        StringBuilder sb = new StringBuilder();
        sb.append(s);
        sb.reverse();
        return sb.toString();
    }

    public WordFilter(String[] words) {
        for(int i = 0; i < words.length; ++i){
            trie.insert(words[i], i, false);
            trie.insert(reverseString(words[i]), i, true);
        }
    }

    public int f(String pref, String suff) {
      Set<Integer> wordWithPrefIdx = trie.startsWith(pref, false);
      Set<Integer> wordWithSuffIdx = trie.startsWith(reverseString(suff), true);

      if(wordWithPrefIdx == null || wordWithSuffIdx == null) return -1;

      wordWithPrefIdx.retainAll(wordWithSuffIdx);

      return wordWithPrefIdx.isEmpty() ? -1 : Collections.max(wordWithPrefIdx);
    }
}

JavaScript

class TrieNode {
  constructor() {
    this.children = Array(26).fill(null);
    this.isWord = false;
    this.wordIdx = new Set();
  }
}

class Trie {
  constructor() {
    this.suffixRoot = new TrieNode();
    this.prefixRoot = new TrieNode();
  }

  insert(word, idx, isSuffixRoot) {
    const iter = isSuffixRoot ? this.suffixRoot : this.prefixRoot;
    for (let i = 0; i < word.length; ++i) {
      const letter = word.charAt(i);
      const charCode = letter.charCodeAt(0) - 97;
      if (iter.children[charCode] === null) {
        iter.children[charCode] = new TrieNode();
      }
      iter = iter.children[charCode];
      iter.wordIdx.add(idx);
    }
    iter.isWord = true;
    iter.wordIdx.add(idx);
  }

  searchPrefix(prefix, isSuffixRoot) {
    const iter = isSuffixRoot ? this.suffixRoot : this.prefixRoot;
    for (let i = 0; i < prefix.length; ++i) {
      const letter = prefix.charAt(i);
      const charCode = letter.charCodeAt(0) - 97;
      if (iter.children[charCode] === null) {
        return null;
      }
      iter = iter.children[charCode];
    }
    return iter;
  }

  startsWith(prefix, suffixLookup) {
    const node = this.searchPrefix(prefix, suffixLookup);
    return node == null ? null : node.wordIdx;
  }
}

class WordFilter {
  constructor(words) {
    this.trie = new Trie();

    for (let i = 0; i < words.length; ++i) {
      this.trie.insert(words[i], i, false);
      this.trie.insert(words[i].split("").reverse().join(""), i, true);
    }
  }

  f(pref, suff) {
    const wordWithPrefIdx = this.trie.startsWith(pref, false);
    const wordWithSuffIdx = this.trie.startsWith(
      suff.split("").reverse().join(""),
      true
    );

    if (wordWithPrefIdx === null || wordWithSuffIdx === null) return -1;

    const intersection = new Set(
      [...wordWithPrefIdx].filter((x) => wordWithSuffIdx.has(x))
    );

    return intersection.size === 0 ? -1 : Math.max(...intersection);
  }
}

Python

class TrieNode:
    def __init__(self):
        self.children = [None] * 26
        self.isWord = False
        self.wordIdx = set()

class Trie:
    def __init__(self):
        self.suffixRoot = TrieNode()
        self.prefixRoot = TrieNode()

    def insert(self, word, idx, isSuffixRoot):
        iter = self.suffixRoot if isSuffixRoot else self.prefixRoot
        for letter in word:
            charCode = ord(letter) - 97
            if iter.children[charCode] is None:
                iter.children[charCode] = TrieNode()
            iter = iter.children[charCode]
            iter.wordIdx.add(idx)
        iter.isWord = True
        iter.wordIdx.add(idx)

    def search_prefix(self, prefix, isSuffixRoot):
        iter = self.suffixRoot if isSuffixRoot else self.prefixRoot
        for letter in prefix:
            charCode = ord(letter) - 97
            if iter.children[charCode] is None:
                return None
            iter = iter.children[charCode]
        return iter

    def starts_with(self, prefix, suffixLookup):
        node = self.search_prefix(prefix, suffixLookup)
        return None if node is None else node.wordIdx

class WordFilter:
    def __init__(self, words):
        self.trie = Trie()
        for i, word in enumerate(words):
            self.trie.insert(word, i, False)
            self.trie.insert(word[::-1], i, True)

    def f(self, pref, suff):
        word_with_pref_idx = self.trie.starts_with(pref, False)
        word_with_suff_idx = self.trie.starts_with(suff[::-1], True)

        if word_with_pref_idx is None or word_with_suff_idx is None:
            return -1

        intersection = word_with_pref_idx.intersection(word_with_suff_idx)

        return max(intersection) if intersection else -1

C++

class TrieNode {
public:
    TrieNode* children[26];
    bool isWord;
    unordered_set<int> wordIdx;

    TrieNode() {
        for (int i = 0; i < 26; ++i) {
            children[i] = nullptr;
        }
        isWord = false;
    }
};

class Trie {
public:
    TrieNode* suffixRoot;
    TrieNode* prefixRoot;

    Trie() {
        suffixRoot = new TrieNode();
        prefixRoot = new TrieNode();
    }

    void insert(string word, int idx, bool isSuffixRoot) {
        TrieNode* iter = isSuffixRoot ? suffixRoot : prefixRoot;
        for (char letter : word) {
            int charCode = letter - 'a';
            if (iter->children[charCode] == nullptr) {
                iter->children[charCode] = new TrieNode();
            }
            iter = iter->children[charCode];
            iter->wordIdx.insert(idx);
        }
        iter->isWord = true;
        iter->wordIdx.insert(idx);
    }

    TrieNode* searchPrefix(string prefix, bool isSuffixRoot) {
        TrieNode* iter = isSuffixRoot ? suffixRoot : prefixRoot;
        for (char letter : prefix) {
            int charCode = letter - 'a';
            if (iter->children[charCode] == nullptr) {
                return nullptr;
            }
            iter = iter->children[charCode];
        }
        return iter;
    }

    unordered_set<int>* startsWith(string prefix, bool suffixLookup) {
        TrieNode* node = searchPrefix(prefix, suffixLookup);
        return node == nullptr ? nullptr : &(node->wordIdx);
    }
};

class WordFilter {
public:
    Trie trie;

    WordFilter(vector<string>& words) {
        for (int i = 0; i < words.size(); ++i) {
            trie.insert(words[i], i, false);
            reverse(words[i].begin(), words[i].end());
            trie.insert(words[i], i, true);
        }
    }

    int f(string pref, string suff) {
        unordered_set<int>* wordWithPrefIdx = trie.startsWith(pref, false);
        reverse(suff.begin(), suff.end());
        unordered_set<int>* wordWithSuffIdx = trie.startsWith(suff, true);

        if (wordWithPrefIdx == nullptr || wordWithSuffIdx == nullptr) return -1;

        vector<int> intersection;
        for (int idx : *wordWithPrefIdx) {
            if (wordWithSuffIdx->count(idx) > 0) {
                intersection.push_back(idx);
            }
        }

        return intersection.empty() ? -1 : *max_element(intersection.begin(), intersection.end());
    }
};

Space complexity is the size of the tries O(nk) number of words and the average (or maximum) length of words, while we have O(nk + Q(n + k)) for time complexity since it takes O(nk) to build the tree and then for each query Q we would have the time required to search for k-length word and time required to traverse the set of words containing certain prefix/suffix leading to O(n + k) (the set can containing all words that have been inserted in the trie) total time complexity.

To save time on read/queries since they will be far more frequent that write/create operation (only once) we can spend more time on the initialization. If we have for example word "cat", we could insert it together with all of its suffixes separated by special (non-lowercase English letter '{'), so we would insert "{cat", "t{cat", "at{cat", "cat{cat", then on query we would query for "t{ca" if the required suffix is "t" and prefix "ca". This leads to a k time complexity for each query Q, while we need more space and time to build the trie in this case O(nk^2) for each word we spend k^2 (k average length of the word) since we're inserting each suffix and there are k (length of word) of those. We use the weight to overwrite the last index of the word that has same path in the trie as one of the previous words, so we can always return the maximum index.