力扣 745. 前缀和后缀搜索

题目来源：https://leetcode.cn/problems/prefix-and-suffix-search/

大致题意：
设计一个类，初始化时给定一个字符串数组作为字典，该类包含一个方法：

• f(string pref, string suff) 返回词典中具有前缀 prefix 和后缀 suff 的单词的下标。如果存在不止一个满足要求的下标，返回其中 最大的下标 。如果不存在这样的单词，返回 -1 。

思路

1. 使用前缀树存下所有前缀和对应的索引，以及字典原字符串倒置后的前缀树，这里称其为后缀树
2. 对于方法 f(string pref, string suff)，判断前缀树中是否存在 pref，若存在取出该前缀对应的索引集合，同理，取出后缀树的索引集合
3. 若其中索引集合有一个为空，那么显然字典中不存在该单词，返回 -1 即可；否则，在两个索引集合中查找单词索引，若找到则返回最大的（这里使用双指针倒序遍历集合，第一个找到的就是最大的），若未找到则返回 -1

代码：

public class WordFilter {
    
    Trie preRoot;   // 前缀树根节点
    Trie sufRoot;   // 后缀树根节点

    public WordFilter(String[] words) {
        preRoot = new Trie();
        sufRoot = new Trie();
        StringBuffer sb;
        // 将给定字典中所有单词插入前缀树和后缀树中
        for (int i = 0; i < words.length; i++) {
            sb = new StringBuffer();
            insert(preRoot, words[i], i);
            sb.append(words[i]);
            insert(sufRoot, sb.reverse().toString(), i);
        }
    }

    /**
     * 判断字典中是否存在前缀为 pref，后缀为 suff 的单词
     * @param pref
     * @param suff
     * @return
     */
    public int f(String pref, String suff) {
        List<Integer> preList = search(preRoot, pref);
        StringBuffer sb = new StringBuffer();
        sb.append(suff).reverse();
        List<Integer> sufList = search(sufRoot, sb.toString());
        // 若有一个索引集合为空，则不存在该单词
        if (preList == null || sufList == null) {
            return -1;
        }
        int idx1 = preList.size() - 1;
        int idx2 = sufList.size() - 1;
        // 双指针倒序遍历
        while (idx1 >= 0 && idx2 >= 0) {
            int preListIdx = preList.get(idx1);
            int sufListIdx = sufList.get(idx2);
            // 找到，直接返回
            if (preListIdx == sufListIdx) {
                return preListIdx;
            } else if (preListIdx > sufListIdx) {
                idx1--;
            } else {
                idx2--;
            }
        }
        // 未找到，返回 -1
        return -1;
    }

    /**
     * 讲单词插入给定前缀树中
     * @param root
     * @param word
     * @param index
     */
    public void insert(Trie root, String word, int index) {
        Trie trie = root;
        for (int i = 0; i < word.length(); i++) {
            int idx = word.charAt(i) - 'a';
            if (trie.children[idx] == null) {
                trie.children[idx] = new Trie();
            }
            trie = trie.children[idx];
            // 插入索引
            trie.indexList.add(index);
        }
    }

    /**
     * 在给定前缀树中搜索给定前缀，并返回索引集合
     * @param root
     * @param prefix
     * @return
     */
    public List<Integer> search(Trie root, String prefix) {
        Trie trie = root;
        for (int i = 0; i < prefix.length(); i++) {
            int idx = prefix.charAt(i) - 'a';
            if (trie.children[idx] == null) {
                return null;
            }
            trie = trie.children[idx];
        }
        return trie.indexList;
    }

    // 前缀树类
    class Trie {
        Trie[] children;
        List<Integer> indexList;    // 用来存含有当前前缀的字符串索引集合

        public Trie() {
            children = new Trie[26];
            indexList = new ArrayList<>();
        }
    }
}