lc-go/solutions/2/q212/solution.go
2026-01-07 18:04:39 +09:00

89 lines
1.7 KiB
Go

package q212
type TrieNode struct {
children [26]*TrieNode
word *string
}
func trieIdx(char byte) int8 {
if char < 'a' || char > 'z' {
return -1
}
return int8(char - 'a')
}
func allocator(bulkSize int) func() *TrieNode {
var v []TrieNode
p := 0
return func() *TrieNode {
if v == nil || p == bulkSize {
v = make([]TrieNode, bulkSize)
p = 0
}
p++
return &v[p-1]
}
}
func buildTrie(words []string) *TrieNode {
alloc := allocator(128)
root := alloc()
for _, word := range words {
p := root
for i := range len(word) {
idx := trieIdx(word[i])
if idx < 0 {
continue
}
if p.children[idx] == nil {
p.children[idx] = alloc()
}
p = p.children[idx]
}
p.word = &word
}
return root
}
type Coord [2]int8
func find(board [][]byte, coord Coord, trie *TrieNode, found *[]string) {
if coord[0] < 0 || coord[1] < 0 ||
int(coord[0]) >= len(board) || int(coord[1]) >= len(board[0]) {
return
}
c := board[coord[0]][coord[1]]
idx := trieIdx(c)
if idx == -1 || trie.children[idx] == nil {
return
}
// Mark visited
board[coord[0]][coord[1]] = '.'
trie = trie.children[idx]
if trie.word != nil {
*found = append(*found, *trie.word)
trie.word = nil // dedup
}
find(board, Coord{coord[0] + 1, coord[1]}, trie, found)
find(board, Coord{coord[0] - 1, coord[1]}, trie, found)
find(board, Coord{coord[0], coord[1] + 1}, trie, found)
find(board, Coord{coord[0], coord[1] - 1}, trie, found)
board[coord[0]][coord[1]] = c
}
func findWords(board [][]byte, words []string) []string {
trie := buildTrie(words)
found := []string{}
for row := range board {
for col := range board[0] {
find(board, Coord{int8(row), int8(col)}, trie, &found)
}
}
return found
}
var _ = findWords