1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91
| class WordFilter {
private static class Trie { private final Trie[] children; private final List<Integer> prefixIndexList; private final List<Integer> suffixIndexList;
public Trie() { this.children = new Trie[26]; this.prefixIndexList = new ArrayList<>(); this.suffixIndexList = new ArrayList<>(); }
public void insert(String word, int wordIndex) { Trie node = this; for (int i = 0; i < word.length(); i++) { int index = word.charAt(i) - 'a'; if (node.children[index] == null) { node.children[index] = new Trie(); } node = node.children[index]; node.prefixIndexList.add(wordIndex); } }
public void reverseInsert(String word, int wordIndex) { Trie node = this; for (int i = word.length() - 1; i >= 0; i--) { int index = word.charAt(i) - 'a'; if (node.children[index] == null) { node.children[index] = new Trie(); } node = node.children[index]; node.suffixIndexList.add(wordIndex); } }
public int search(String pref, String suff) { Trie node = this; for (int i = 0; i < pref.length() && node != null; i++) { int index = pref.charAt(i) - 'a'; node = node.children[index]; } if (node == null || node.prefixIndexList.isEmpty()) { return -1; }
List<Integer> indexListA = node.prefixIndexList;
node = this; for (int i = suff.length() - 1; i >= 0 && node != null; i--) { int index = suff.charAt(i) - 'a'; node = node.children[index]; } if (node == null || node.suffixIndexList.isEmpty()) { return -1; } List<Integer> indexListB = node.suffixIndexList;
int i = indexListA.size() - 1, j = indexListB.size() - 1; while (i >= 0 && j >= 0) { int indexA = indexListA.get(i); int indexB = indexListB.get(j); if (indexA < indexB) { j--; } else if (indexA > indexB) { i--; } else { return indexA; } }
return -1; } }
private final Trie trie;
public WordFilter(String[] words) { this.trie = new Trie(); for (int i = 0; i < words.length; i++) { trie.insert(words[i], i); trie.reverseInsert(words[i], i); } }
public int f(String pref, String suff) { return trie.search(pref, suff); }
}
|
References
745. Prefix and Suffix Search