1
class TrieNode {
2
 public:
3
  vector<TrieNode *> children;
4
  vector<int> wordIndexList;
5

6
  TrieNode() {
7
    children = vector<TrieNode *>(26, NULL);
8
  }
9
};
10

11
class Trie {
12
 public:
13
  // Root is a datamember of trie.
14
  TrieNode *root;
15

16
  // Constructor
17
  Trie() {
18
    // Initialize root.
19
    root = new TrieNode();
20
  }
21

22
  void insert(string &word, int wordIndex) {
23
    auto cur = root;
24
    for (int i = 0; i < word.size(); i++) {
25
      int index = word[i] - 'a';
26
      if (!cur->children[index]) cur->children[index] = new TrieNode();
27
      cur = cur->children[index];
28
      cur->wordIndexList.push_back(wordIndex);
29
    }
30
  }
31

32
  vector<int> searchPrefOrSuff(string &str) {
33
    auto cur = root;
34
    int n = str.size();
35
    for (int i = 0; i < n; i++) {
36
      int index = str[i] - 'a';
37
      if (cur->children[index] == NULL) return {};
38
      cur = cur->children[index];
39
    }
40
    return cur->wordIndexList;
41
  }
42
};
43

44
class WordFilter {
45
 public:
46
  Trie *prefixTrie = new Trie();
47
  Trie *suffixTrie = new Trie();
48
  unordered_map<string, int> mp;
49

50
  WordFilter(vector<string> &words) {
51
    int n = words.size();
52
    for (int i = 0; i < n; i++) {
53
      string word = words[i];
54
      prefixTrie->insert(word, i);
55
      string rev = word;
56

57
      reverse(rev.begin(), rev.end());
58
      suffixTrie->insert(rev, i);
59
    }
60
  }
61

62
  int f(string pref, string suff) {
63
    string key = pref + "-" + suff;
64
    if (mp.find(key) != mp.end()) return mp[key];
65

66
    vector<int> prefList = prefixTrie->searchPrefOrSuff(pref);
67
    reverse(suff.begin(), suff.end());
68
    vector<int> suffList = suffixTrie->searchPrefOrSuff(suff);
69

70
    int n = prefList.size(), m = suffList.size();
71
    int k = n - 1, l = m - 1;
72

73
    while (k >= 0 && l >= 0) {
74
      if (prefList[k] == suffList[l]) return mp[key] = prefList[k];
75
      if (prefList[k] > suffList[l])
76
        k -= 1;
77
      else if (suffList[l] > prefList[k])
78
        l -= 1;
79
    }
80
    return mp[key] = -1;
81
  }
82
};