LeetCode-in-Java

Hard

Design a special dictionary with some words that searchs the words in it by a prefix and a suffix.

Implement the WordFilter class:

Example 1:

Input

[“WordFilter”, “f”]

[[[“apple”]], [“a”, “e”]]

Output: [null, 0]

Explanation:

WordFilter wordFilter = new WordFilter(["apple"]); 
wordFilter.f("a", "e"); // return 0, because the word at index 0 has prefix = "a" and suffix = 'e".

Constraints:

Solution

public class WordFilter {
    private static class TrieNode {
        TrieNode[] children;
        int weight;

        TrieNode() {
            this.children = new TrieNode[27];
            this.weight = 0;
        }
    }

    private TrieNode root = new TrieNode();

    public WordFilter(String[] words) {
        for (int i = 0; i < words.length; i++) {
            String s = words[i];
            for (int j = 0; j < s.length(); j++) {
                insert(s.substring(j) + '{' + s, i);
            }
        }
    }

    public void insert(String wd, int weight) {
        TrieNode node = root;
        for (char ch : wd.toCharArray()) {
            if (node.children[ch - 'a'] == null) {
                node.children[ch - 'a'] = new TrieNode();
            }
            node = node.children[ch - 'a'];
            node.weight = weight;
        }
    }

    public int f(String prefix, String suffix) {
        return startsWith(suffix + '{' + prefix);
    }

    public int startsWith(String pref) {
        TrieNode node = root;
        for (char ch : pref.toCharArray()) {
            if (node.children[ch - 'a'] == null) {
                return -1;
            }
            node = node.children[ch - 'a'];
        }
        return node.weight;
    }
}