745. Prefix and Suffix Search

Given many words, words[i] has weight i.

Design a class WordFilter that supports one function, WordFilter.f(String prefix, String suffix). It will return the word with given prefix and suffix with maximum weight. If no word exists, return -1.

Examples:

Input:
WordFilter(["apple"])
WordFilter.f("a", "e") // returns 0
WordFilter.f("b", "") // returns -1

Note:

1. words has length in range [1, 15000].

2. For each test case, up to words.length queries WordFilter.f may be made.

3. words[i] has length in range [1, 10].

4. prefix, suffix have lengths in range [0, 10].

5. words[i] and prefix, suffix queries consist of lowercase letters only.

class WordFilter {
public:
    WordFilter(vector<string>& words) { // time: O(N * L^2); space: O(N * L^2)
        for (int k = 0; k < words.size(); ++k) {
            for (int i = 0; i <= words[k].length(); ++i) {
                for (int j = 0; j <= words[k].length(); ++j) {
                    m[words[k].substr(0, i) + "#" + words[k].substr(words[k].length() - j)] = k;
                }
            }
        }
    }
    
    int f(string prefix, string suffix) { // time: O(1)
        string key_str = prefix + "#" + suffix;
        return m.count(key_str) != 0 ? m[key_str] : -1;
    }
private:
    unordered_map<string, int> m;
};

class WordFilter {
public:
    WordFilter(vector<string>& words) {
        for (int k = 0; k < words.size(); ++k) {
            for (int i = 0; i <= words[k].length(); ++i) {
                prefix_mp[words[k].substr(0, i)].emplace_back(k);
                suffix_mp[words[k].substr(words[k].length() - i)].emplace_back(k);
            }
        }
    }
    
    int f(string prefix, string suffix) {
        if (!prefix_mp.count(prefix) || !suffix_mp.count(suffix)) return -1;
        vector<int> pre = prefix_mp[prefix], suf = suffix_mp[suffix];
        int i = pre.size() - 1, j = suf.size() - 1;
        while (i >= 0 && j >= 0) {
            if (pre[i] < suf[j]) --j;
            else if (pre[i] > suf[j]) --i;
            else return pre[i];
        }
        return -1;
    }
private:
    unordered_map<string, vector<int> > prefix_mp, suffix_mp;
};

class WordFilter {
public:
    struct TrieNode {
        vector<TrieNode*> next;
        vector<int> words;
        TrieNode() : next(26, nullptr) {}
        ~TrieNode() {
            for (auto& ptr : next) {
                if (ptr) delete ptr;
            }
        }
    };
    
    void build(TrieNode *node, string& word, int k, int weight, bool reverse) {
        if ( (!reverse && k == word.size() || (reverse && k == -1)) ) return;
        int pos = word[k] - 'a';
        if (node->next[pos] == nullptr) node->next[pos] = new TrieNode();
        node->next[pos]->words.push_back(weight);
        build(node->next[pos], word, k + (reverse ? -1 : 1), weight, reverse);
    }
    
    TrieNode* search(TrieNode *node, string& word, int k) {
        if (word.size() == k) return node;
        int pos = word[k] - 'a';
        return !node->next[pos] ? nullptr : search(node->next[pos], word, k + 1);
    }
    
    WordFilter(vector<string>& words) {
        prefixRoot = new TrieNode();
        suffixRoot = new TrieNode();
        for (int i = 0; i < words.size(); i++) {
            prefixRoot->words.push_back(i);
            suffixRoot->words.push_back(i);
            build(prefixRoot, words[i], 0, i, false);
            build(suffixRoot, words[i], words[i].size()-1, i, true);
        }
    }
    
    int f(string prefix, string suffix) {
        TrieNode *p = search(prefixRoot, prefix, 0);
        reverse(suffix.begin(), suffix.end());
        TrieNode *s = search(suffixRoot, suffix, 0);
        if (!p || !s) return -1;
         
        int pI = p->words.size()-1;
        int sI = s->words.size()-1;
         
        while (pI >= 0 && sI >= 0) {
            int pW = p->words[pI];
            int sW = s->words[sI];
            if (pW == sW) return pW;
            if (pW > sW) pI--;
            else sI--;
        }
         
        return -1;
    }
    
private:
    TrieNode *prefixRoot;
    TrieNode *suffixRoot;
};