212. Word Search II

Problem Link

Java Version (Accepted)

class Node {
    Map<String,Node> children;
    boolean isEnd;

    public Node() {

        this.children = new HashMap<String,Node>();
        this.isEnd = false;
    }
}

class Trie {
    Node root;

    public Trie() {
        this.root = new Node();
    }

    public void insert(String word) {
        Node cur = this.root;
        char[] charArray = word.toCharArray();
        for (char c : charArray) {
            String strC = Character.toString(c);
            if (!cur.children.containsKey(strC)) {
                Node newNode = new Node();
                cur.children.put(strC, newNode);
            }
            cur = cur.children.get(strC);
        }
        cur.isEnd = true;
    }
}

class Solution {
    public List<String> findWords(char[][] board, String[] words) {
        Trie trie = new Trie();
        for (String word : words) {
            trie.insert(word);
        }
        int M = board.length;
        int N = board[0].length;
        List<String> res = new ArrayList<String>();
        boolean[][] visited = new boolean[M][N];
        for (int i=0; i<M; i++) {
            for (int j=0; j<N;j++) {
                backtrack(board, i, j, trie.root, visited, "", res);
            }
        }
        return res;
    }

    public void backtrack(char[][] board, int i, int j, Node node, boolean[][] visited, String cur, List<String> res) {
        String curChar = Character.toString(board[i][j]);
        if (!node.children.containsKey(curChar)) {
            return;
        }
        cur += curChar;
        if (node.children.get(curChar).isEnd) {
            res.add(String.valueOf(cur));
            node.children.get(curChar).isEnd = false;
        }
        int M = board.length;
        int N = board[0].length;
        visited[i][j] = true;
        int ni;
        int nj;
        ni = i + 1;
        nj = j;
        if (0 <= ni && ni < M && 0 <= nj && nj < N && !visited[ni][nj])
            backtrack(board, ni, nj, node.children.get(curChar), visited, cur, res);
        ni = i;
        nj = j + 1;
        if (0 <= ni && ni < M && 0 <= nj && nj < N && !visited[ni][nj])
            backtrack(board, ni, nj, node.children.get(curChar), visited, cur, res);
        ni = i - 1;
        nj = j;
        if (0 <= ni && ni < M && 0 <= nj && nj < N && !visited[ni][nj])
            backtrack(board, ni, nj, node.children.get(curChar), visited, cur, res);
        ni = i;
        nj = j - 1;
        if (0 <= ni && ni < M && 0 <= nj && nj < N && !visited[ni][nj])
            backtrack(board, ni, nj, node.children.get(curChar), visited, cur, res);
        cur = cur.substring(0, cur.length() - 1);
        visited[i][j] = false;
    }
}

Python Version (TLE)

class Node:
    def __init__(self):
        self.children = defaultdict(Node)
        self.isEnd = False

class Trie:

    def __init__(self):
        self.root = Node()

    def insert(self, word: str) -> None:
        cur = self.root
        for c in word:
            cur = cur.children[c]
        cur.isEnd = True

class Solution:
    def findWords(self, board: List[List[str]], words: List[str]) -> List[str]:
        trie = Trie()
        for word in words:
            trie.insert(word)
        M, N = len(board), len(board[0])
        res = []
        visited = [[False for _ in range(N)] for _ in range(M)]
        for i in range(M):
            for j in range(N):
                self.backtrack(board, i, j, trie.root, visited, '', res)
        return list(res)

    def backtrack(self, board, i, j, node, visited, cur, res) -> None:
        curChar = board[i][j]
        if curChar not in node.children:
            return 
        cur += curChar
        if node.children[curChar].isEnd:
            res.append(cur[::])
            node.children[curChar].isEnd = False
        dirs = {(1,0), (0,1), (-1,0), (0,-1)}
        M, N = len(board), len(board[0])
        visited[i][j] = True
        for d in dirs:
            ni, nj = i + d[0], j + d[1]
            if 0 <= ni < M and 0 <= nj < N and not visited[ni][nj]:
                self.backtrack(board, ni, nj, node.children[curChar], visited, cur, res)
        # cur = cur[:-1]
        visited[i][j] = False