LeetCode: 212. Word Search II

단순 구현 문제로 생각했는데, 시간 초과가 발생한다.

메모리를 더 써서 시간을 줄여야 한다.

class Solution:
    def findWords(self, board: List[List[str]], words: List[str]) -> List[str]:
        m, n = len(board), len(board[0])
        ret = []
        for word in words:
            found_flag = False
            for i in range(m):
                for j in range(n):
                    if board[i][j] == word[0]:
                        visited = [[False] * n for _ in range(m)]
                        visited[i][j] = True
                        if self.dfs(board, visited, i, j, word[1:]):
                            ret.append(word)
                            found_flag = True
                            break
                if found_flag:
                    break
        return ret
    
    
    def dfs(self, board, visited, y, x, word):
        if not word:
            return True
        dy = [0, 1, 0, -1]
        dx = [1, 0, -1, 0]
        for i in range(4):
            next_y = y + dy[i]
            next_x = x + dx[i]
            if 0 <= next_y <= len(board) - 1 and 0 <= next_x <= len(board[0]) - 1:
                if board[next_y][next_x] == word[0] and not visited[next_y][next_x]:
                    visited[next_y][next_x] = True
                    if self.dfs(board, visited, next_y, next_x, word[1:]):
                        return True
                    visited[next_y][next_x] = False
        return False

 

TRIE 구조를 사용해야 한다.

TRIE는 아래 글을 참고 바란다.

https://focalpoint.tistory.com/214

LeetCode: 208. Implement Trie (Prefix Tree) - 시바견의 끄적임

이 문제의 알고리즘 자체는 간단하다.

보드의 모든 좌표에 대해서 TRIE를 체크하면 된다.

시간복잡도는? O((단어 갯수 * 단어 최대 길이) * (보드 세로 * 보드 가로))

class TrieNode:
    def __init__(self):
        self.children = {}
        self.endFlag = False

        
class Trie:
    def __init__(self):
        self.root = TrieNode()
        
    def insert(self, word: str) -> None:
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
        node.endFlag = True

    def search(self, word: str) -> bool:
        node = self.root
        for char in word:
            if char not in node.children:
                return False
            node = node.children[char]
        return node.endFlag
        
    def startsWith(self, prefix: str) -> bool:
        node = self.root
        for char in prefix:
            if char not in node.children:
                return False
            node = node.children[char]
        return True
    
    def delete(self, word):
        def recursive(node, word, i):
            if i == len(word):
                if not node.endFlag:
                    return False
                node.endFlag = False
                return (len(node.children) == 0)
            if word[i] not in node.children:
                return False
            need_delete = recursive(node.children[word[i]], word, i+1)
            if need_delete:
                node.children.pop(word[i])
                return (len(node.children) == 0)
            return False
        recursive(self.root, word, 0)

        
class Solution:
    def findWords(self, board: List[List[str]], words: List[str]) -> List[str]:
        self.trie = Trie()
        self.root = self.trie.root
        for word in words:
            self.trie.insert(word)
        self.ret = []
        m, n = len(board), len(board[0])
        for y in range(m):
            for x in range(n):
                self.dfs(self.root, board, y, x, '')
        return self.ret
                
    def dfs(self, node, board, y, x, path):
        if node.endFlag:
            self.ret.append(path)
            # avoid appending duplicate words
            node.endFlag = False
        # invalid coordinate
        if y < 0 or y >= len(board) or x < 0 or x >= len(board[0]):
            return
        char = board[y][x]
        node = node.children.get(char)
        if not node:
            return
        board[y][x] = '#'
        self.dfs(node, board, y, x+1, path+char)
        self.dfs(node, board, y+1, x, path+char)
        self.dfs(node, board, y, x-1, path+char)
        self.dfs(node, board, y-1, x, path+char)
        board[y][x] = char

Success이긴하나 조금 느리다... 개선할 여지가 있을까?

1) Trie 구조를 단순화한다 -> 메모리 개선

2) 검색이 완료된 Trie를 제거한다 -> 시간 개선

class Solution:
    def findWords(self, board, words):
        self.ret = []
        self.trie = {}
        for word in words:
            t = self.trie
            for c in word:
                if c not in t:
                    t[c] = {}
                t = t[c]
            t['#'] = '#'
        
        m, n = len(board), len(board[0])
        for y in range(m):
            for x in range(n):
                self.dfs(board, self.trie, y, x, '')
        return self.ret
    
    def dfs(self, board, node, y, x, word):
        if '#' in node:
            del node['#']
            self.ret.append(word)
        
        if 0 > y or y >= len(board) or 0 > x or x >= len(board[0]):
            return
        
        char = board[y][x]
        if char not in node:
            return
        
        next_node = node[char]
        board[y][x] = '@'
        dy, dx = [0, 1, 0, -1], [1, 0, -1, 0]
        for d in range(4):
            next_y, next_x = y + dy[d], x + dx[d]
            self.dfs(board, next_node, next_y, next_x, word+char)
        board[y][x] = char
        if not next_node:
            del node[char]