Word Search II

Given a 2D board and a list of words from the dictionary, find all words in the board.

Each word must be constructed from letters of sequentially adjacent cell, where "adjacent" cells are those horizontally or vertically neighboring. The same letter cell may not be used more than once in a word.

 

Example:

Input: 
board = [
  ['o','a','a','n'],
  ['e','t','a','e'],
  ['i','h','k','r'],
  ['i','f','l','v']
]
words = ["oath","pea","eat","rain"]

Output: ["eat","oath"]

 

Note:

  1. All inputs are consist of lowercase letters a-z.
  2. The values of words are distinct.

Solution:

class Solution {
    static class TrieNode {
        List<TrieNode> children;
        char c;
        
        public TrieNode(char c) {
            this.children = new ArrayList<>();
            this.c = c;
        }
    }
    
    static class Trie {
        TrieNode root;
        
        public Trie() {
            root = new TrieNode('\0');
        }
        
        public void insert(String s) {
            char[] arr = s.toCharArray();
            TrieNode curr = root;
            for (char c : arr) {
                boolean found = false;
                for (TrieNode child : curr.children) {
                    if (child.c == c) {
                        curr = child;
                        found = true;
                        break;
                    }
                }
                if (!found) {
                    TrieNode node = new TrieNode(c);
                    curr.children.add(node);
                    curr = node;
                }
            }
        }
        
        public boolean isPrefix(String str) {
            List<TrieNode> nodes = root.children;
            for (char c : str.toCharArray()) {
                boolean found = false;
                for (TrieNode node : nodes) {
                    if (node.c == c) {
                        nodes = node.children;
                        found = true;
                        break;
                    }
                }
                if (!found) return false;
            }
            return true;
        }
    }
    
    public List<String> findWords(char[][] board, String[] words) {
        Set<String> dict = new HashSet();
        Trie trie = new Trie();
        for (String s : words) {
            trie.insert(s);
            dict.add(s);
        }
        Set<String> result = new HashSet();
        int m = board.length;
        if (m == 0) return new ArrayList();
        int n = board[0].length;
        if (n == 0) return new ArrayList();
        for (int i = 0; i < m; i ++) {
            for (int j = 0; j < n; j ++) {
                if (dict.size() == 0) break;
                StringBuilder sb = new StringBuilder();
                sb.append(board[i][j]);
                boolean[][] visited = new boolean[m][n];
                visited[i][j] = true;
                dfs(board, dict, trie, visited, sb, result, i, j);
            }
        }
        List<String> r = new ArrayList(result);
        return r;
    }
    
    private void dfs(char[][] board, Set<String> dict, Trie trie, boolean[][] visited, StringBuilder sb, Set<String> result, int i, int j) {
        boolean lengthTooLong = true;
        for (String s : dict.toArray(new String[dict.size()])) {
            if (s.length() > sb.length()) lengthTooLong = false;
            if (sb.toString().equals(s)) {
                dict.remove(s);
                result.add(s);
            }
        }
        if (lengthTooLong) return;
        int[] dx = new int[]{-1,0,1,0};
        int[] dy = new int[]{0,1,0,-1};
        int m = board.length;
        int n = board[0].length;
        for (int k = 0; k < 4; k ++) {
            int nx = i + dx[k];
            int ny = j + dy[k];
            if (nx >= 0 && nx < m && ny >= 0 && ny < n && !visited[nx][ny]) {
                visited[nx][ny] = true;
                sb.append(board[nx][ny]);
                if (trie.isPrefix(sb.toString())) {
                    dfs(board, dict, trie, visited, sb, result, nx, ny);
                }
                sb.setLength(sb.length() - 1);
                visited[nx][ny] = false;
            }
        }
    }
}