LeetCode - 解题笔记 - 212 - Word Search II

Solution 1

针对 0079. Word Search 的扩展，针对一个词的判定扩展到整个列表，并找到所有的词。

根据题目给出的输入大小，直接在单个词DFS上套一个遍历是肯定会超时的，因此需要简化单词的查询逻辑：前缀树。DFS的深入遍历过程可以直接关联前缀树的向下探查和向上递归的过程。

本题基于 0208. Implement Trie (Prefix Tree) 的实现改进，增加了一些额外的状态读写接口就可以直接用了。

• 时间复杂度：$O(MN\times 3^L)$，其中$M$和$N$时输入地图的高度和宽度，$L$为字符串长度。搜索树的分支数最多为3，每一次递归都直接对应了一次前缀树的探查，状态判断为常数复杂度
• 空间复杂度：$O(MN+LR)$，其中$M$和$N$时输入地图的高度和宽度，$L$为字符串长度，$R$为前缀树基数。每个遍历位置的状态保存以及前缀树保存消耗

class Trie {
private:
    vector<Trie*> next;
    const int R = 26;
    bool isEnd;
    string target;
    
public:
    Trie() {
        next = vector<Trie*>(this->R);
        isEnd = false;
    }
    
    void insert(string word) {
        auto node = this;
        for (char c: word) {
            if (node->next[c - 'a'] == nullptr) {
                // 不存在，直接插入
                node->next[c - 'a'] = new Trie();
            }
            node = node->next[c - 'a']; // 向后缀查询
        }
        
        node->isEnd = true;
        node->target = word;
    }
    
    bool search(string word) {
        auto node = this->searchPrefix(word);
        
        if (node != nullptr && node->isEnd) {
            return true;
        } else {
            return false;
        }
    }
    
    bool startsWith(string prefix) {
        auto node = this->searchPrefix(prefix);
        
        if (node != nullptr) {
            return true;
        } else {
            return false;
        }
    }
    
    bool existsNext(char c) {
        return this->next[c - 'a'] != nullptr;
    }
    
    Trie* getNext(char c) {
        return this->next[c - 'a'];
    }
    
    bool checkEnd() {
        return this->isEnd;
    }
    
    void setEnd(bool newState) {
        this->isEnd = newState;
    }
    
    string getTarget() {
        return this->target;
    }
    
private:
    Trie* searchPrefix(string prefix) {
        auto node = this;
        for (char c: prefix) {
            if (node->next[c - 'a'] == nullptr) {
                // 不存在，搜索失败
                return nullptr;
            }
            node = node->next[c - 'a']; // 向后缀查询
        }
        
        return node;
    }
};

class Solution {
public:
    vector<string> findWords(vector<vector<char>>& board, vector<string>& words) {
        vector<string> ans;
        Trie * trie = new Trie();
        for (auto word: words) {
            trie->insert(word);
        }
        
        vector<vector<bool>> visited(board.size(), vector<bool>(board[0].size()));
        for (int i = 0; i < board.size(); ++i) {
            for (int j = 0; j < board[0].size(); ++j) {
                if (trie->existsNext(board[i][j])) {
                    // 基础剪枝：只查询可能存在的词汇
                    this->dfs(board, visited, i, j, trie->getNext(board[i][j]), ans);
                }
            }
        }
        
        return ans;
    }
    
private:
    vector<pair<int, int>> dirs{{-1, 0}, {1, 0}, {0, -1}, {0, 1}};
    
    void dfs(vector<vector<char>> & board, vector<vector<bool>> & visited, int i, int j, Trie * trie, vector<string> & ans) {
        if (trie->checkEnd()) {
            // cout << board[i][j] << " " << i << " " << j << " Found: " << trie->getTarget() << endl;
            ans.push_back(trie->getTarget());
            // return;
            // 考虑到前缀关系：不能直接返回要继续往下探，同时要去掉这一个信息
            trie->setEnd(false);
        }
        
        // cout << board[i][j] << " " << i << " " << j << "->";
        
        visited[i][j] = true;
        
        for (auto & dir: dirs) {
            int iNext = i + dir.first, jNext = j + dir.second;
            if (
                iNext >= 0 && iNext < board.size() &&
                jNext >= 0 && jNext < board[0].size() && 
                !visited[iNext][jNext] &&
                trie->existsNext(board[iNext][jNext])
               ) {
                // cout << board[iNext][jNext] << " " << iNext << " " << jNext << endl;
                this->dfs(board, visited, iNext, jNext, trie->getNext(board[iNext][jNext]), ans);
            }
        }
        
        visited[i][j] = false;
    }
};

Solution 2

Solution 1的Python实现，参考了Python dfs+Trie+pruning solution(2022 not TLE) - LeetCode Discuss的办法做了更多的剪枝。在每个节点记录尚未检查的前缀的个数，并在dfs期间更新之。如果出现了0的情况即所有的后续前缀均已经检查完毕而完全跳过后续搜索，实现更小的搜索空间。

class Trie:
    
    R = 26

    def __init__(self):
        self.next = collections.defaultdict(Trie)
        self.is_end = False
        self.target = None
        self.num_unchecked = 0 # 额外的剪枝：尚未检查的prefix数量

    def insert(self, word: str) -> None:
        node = self
        for c in word:
            if c not in node.next:
                node.next[c] = Trie()
                
            node = node.next[c]
            node.num_unchecked += 1
            
        node.is_end = True
        node.target = word

    def search(self, word: str) -> bool:
        node = self.__search_prefix(word)
        
        if node is not None and node.is_end: return True
        else: return False

    def startsWith(self, prefix: str) -> bool:
        node = self.__search_prefix(prefix)
        
        if node is not None: return True
        else: return False
        
        
    def __search_prefix(self, prefix: str) -> Optional[Any]:
        node = self
        for c in prefix:
            if c not in node.next:
                return None
            
            node = node.next[c]
            
        return node

    
class Solution:
    def findWords(self, board: List[List[str]], words: List[str]) -> List[str]:
        ans = list()
        
        dirs = [(-1, 0), (1, 0), (0, -1), (0, 1)]
        visited = [[False] * len(board[0]) for i in range(len(board))]
        
        #########################################################
        def dfs(i: int, j: int, node: Trie) -> None: 
            # print(i, j, board[i][j], end=' -> ')
            if node.num_unchecked == 0: return # 额外剪枝
            
            nonlocal ans, visited, dirs
            prev_num_ans = len(ans)
            
            if node.is_end:
                ans.append(node.target)
                node.is_end = False
                # print("Found: ", node.target)
            
            visited[i][j] = True
            
            for ioffset, joffset in dirs:
                iNext, jNext = i + ioffset, j + joffset
                if (0 <= iNext < len(board)) and (0 <= jNext < len(board[0])) and (not visited[iNext][jNext]) and (board[iNext][jNext] in node.next):
                    # print(iNext, jNext, board[iNext][jNext])
                    dfs(iNext, jNext, node.next[board[iNext][jNext]])
            
            visited[i][j] = False
            
            num_ans = len(ans)
            node.num_unchecked -= (num_ans - prev_num_ans)
        #############################################################
        
        trie = Trie()
        for word in words:
            trie.insert(word)
        
        for i in range(len(board)):
            for j in range(len(board[0])):
                if board[i][j] in trie.next:
                    # print("Start Checking: ", i, j, board[i][j])
                    dfs(i, j, trie.next[board[i][j]])
        
        return ans