leetcode-go/leetcode/hard/212.word-search-ii.go

func findWords(board [][]byte, words []string) (ans []string) {
	trie := &TrieNode{}
	for i := range words {
		trie.Insert(words[i])
	}
	for y := range board {
		for x := range board[y] {
			DFS(board, trie, trie, x, y, &ans, []byte{})
		}
	}
	return
}

func DFS(board [][]byte, trie, curr *TrieNode, x, y int, ans *[]string, prev []byte) {
	ch := board[y][x]
	if ch == '*' { // Visited
		return
	}
	curr = curr.OneStepSearch(ch)
	if curr == nil {
		return
	}
	path := make([]byte, len(prev))
	copy(path, prev)
	path = append(path, ch)
	if curr.IsKey {
		*ans = append(*ans, string(path))
		trie.Delete(path)
	}
	board[y][x] = '*' // Mark (x, y) as visited
	if 0 < x {
		DFS(board, trie, curr, x-1, y, ans, path)
	}
	if 0 < y {
		DFS(board, trie, curr, x, y-1, ans, path)
	}
	if x < len(board[0])-1 {
		DFS(board, trie, curr, x+1, y, ans, path)
	}
	if y < len(board)-1 {
		DFS(board, trie, curr, x, y+1, ans, path)
	}
	board[y][x] = ch
}

type TrieNode struct {
	Cnt      int
	IsKey    bool
	Children *[26]TrieNode
}

func (root *TrieNode) Insert(word string) {
	for i := range word {
		if root.Children == nil {
			root.Children = &[26]TrieNode{}
		}
		root = &root.Children[int(word[i]-'a')]
		root.Cnt++
	}
	root.IsKey = true
}

func (root *TrieNode) Delete(word []byte) {
	for i := range word {
		root = &root.Children[int(word[i]-'a')]
		root.Cnt--
	}
	root.IsKey = false
}

func (root *TrieNode) OneStepSearch(ch byte) *TrieNode {
	if root.Children == nil {
		return nil
	}
	root = &root.Children[int(ch-'a')]
	if root.Cnt == 0 {
		return nil
	}
	return root
}