面试题 17.17. 多次搜索

面试题 17.17. 多次搜索

https://leetcode-cn.com/problems/multi-search-lcci/

/**
 *
 * 执行用时： 37 ms
 * 内存消耗： 59.9 MB
 */
class Solution {
    public int[][] multiSearch(String big, String[] smalls) {
        if (big == null || smalls == null) {
            return null;
        }
        Trie trie = new Trie();
        trie.insertWords(smalls);
        return trie.searchBigWord(big);
    }
}


class Trie {
    private int position;
    private int wordSize;
    private Trie[] children;
    private boolean end;

    public Trie() {
        position = 0;
        children = new Trie[26];
    }

    public void insertWords(String[] smallWords) {
        wordSize = smallWords.length;

        for (int i = 0; i < smallWords.length; i++) {
            insertPer(smallWords[i], i);
        }
    }

    public void insertPer(String word, int position) {
        Trie node = this;

        for (int i = 0; i < word.length(); i++) {
            if (node.children[word.charAt(i) - 'a'] == null) {
                node.children[word.charAt(i) - 'a'] = new Trie();
            }
            node = node.children[word.charAt(i) - 'a'];
        }
        node.position = position;
        node.end = true;
    }

    public int[][] searchBigWord(String bigWord) {
        List<List<Integer>> temp = new ArrayList<>(this.wordSize);

        for (int i = 0; i < wordSize; i++) {
            temp.add(new ArrayList<Integer>());
        }

        for (int i = 0; i < bigWord.length(); i++) {
            String subWord = bigWord.substring(i);
            Trie node = this;
            for (int j = 0; j < subWord.length(); j++) {
                if (node.children[subWord.charAt(j) - 'a'] == null) {
                    break;
                }
                node = node.children[subWord.charAt(j) - 'a'];
                if (node.end) {
                    temp.get(node.position).add(i);
                }
            }
        }

        int[][] res = new int[temp.size()][];
        for (int i = 0; i < temp.size(); i++) {
            List<Integer> subList = temp.get(i);
            res[i] = new int[subList.size()];
            for (int j = 0; j < subList.size(); j++) {
                res[i][j] = subList.get(j);
            }
        }
        return res;
    }
}