[Java] Trie and Backtracking


  • 1
    A
    
    class TrieNode {
        char c;
        Map<Character, TrieNode> children = new HashMap<>();
        boolean isLeaf;
        public TrieNode() {
        }
        public TrieNode(char c) {
            this.c = c;
        }
    }
    class Trie {
    
        TrieNode root;
        public Trie() {
            root = new TrieNode();
        }
        
        public void addWord(String word) {
    
            if (word == null)
                return;
            TrieNode t = root;
            for (int i = 0; i < word.length(); ++i) {
                char c = word.charAt(i);
                if (!t.children.containsKey(c)) {
                    t.children.put(c, new TrieNode(c));
                }
                t = t.children.get(c);
            }
            t.isLeaf = true;
        }
        
        public List<String> searchPrefix(String word) {
    
            List<String> result = new ArrayList<>();
            TrieNode lastNode = searchNode(word);
            if (lastNode == null) {
                return result;
            }
            dfsSearch(lastNode, result, new StringBuffer(word));
            return result;
        }
    
        private TrieNode searchNode(String str){
            TrieNode t = root;
            for(int i=0; i<str.length(); i++){
                char c = str.charAt(i);
                if(t.children.containsKey(c)){
                    t = t.children.get(c);
                }else{
                    return null;
                }
            }
    
            return t;
        }
    
        // Add in all words start with prefix
        private void dfsSearch(TrieNode parent, List<String> result, StringBuffer prefix) {
    
            if (parent.isLeaf) {
                result.add(prefix.toString());
            }
            for (char c: parent.children.keySet()) {
                prefix.append(c);
                dfsSearch(parent.children.get(c), result, prefix);
                prefix.setLength(prefix.length() - 1);
            }
            return;
        }
    }
    
    // Solution
    public class Solution {
        public List<List<String>> wordSquares(String[] words) {
            List<List<String>> list = new ArrayList<>();
            int wordsLen = words.length;
            int squareLen = words[0].length();
            Trie trie = new Trie();
            for (String word: words) {
                trie.addWord(word);
            }
            for (int i = 0; i < wordsLen; ++i) {
                String[] tmp = new String[squareLen];
                tmp[0] = words[i];
                DFSHelper(1, squareLen, tmp, trie, list);
            }
            return list;
        }
    
        public void DFSHelper(int index, int squareLen, String[] tmp, Trie trie, List<List<String>> list) {
    
            if (index == squareLen) {
                list.add(new ArrayList<>(Arrays.asList(tmp)));
                return;
            }
            // General the prefix we need to start with
            StringBuffer prefix = new StringBuffer();
            for (int i = 0; i < index; ++i) {
                prefix.append(tmp[i].charAt(index));
            }
            List<String> wordsWithPrefix  = trie.searchPrefix(prefix.toString());
            if (wordsWithPrefix  == null || wordsWithPrefix.isEmpty())
                return;
            for (String word: wordsWithPrefix) {
                tmp[index] = word;
                DFSHelper(index + 1, squareLen, tmp, trie, list);
            }
        }
    }
    

  • 0
    S

    Brilliant code man. very well written. So easy to understand. But would appreciate some more comments.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.