字典树，前缀树

Posted on Edited on In Word count in article: 6.6k Reading time ≈ 6 mins.
go的字典树实现，前缀树，以及相关问题

trie 字典树实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
type Trie struct {
    isEnd       bool
    children    [26]*Trie
}


/** Initialize your data structure here. */
func Constructor() Trie {
    return Trie{}
}


/** Inserts a word into the trie. */
func (this *Trie) Insert(word string)  {
    root := this
    // 每个单词字符处理，是nil的话就新建
    for i := 0; i < len(word); i++ {
        if root.children[word[i] - 'a'] == nil {
            root.children[word[i] - 'a'] = &Trie{}
        }
        // 不是nil就更新下一节点去遍历
        root = root.children[word[i] - 'a']
    }
    // 最后一个节点是单词结尾
    root.isEnd = true
}


/** Returns if the word is in the trie. */
func (this *Trie) Search(word string) bool {
    root := this
    // 搜索与插入类似，是nil的话直接返回false
    for i := 0; i < len(word); i++ {
        if root.children[word[i] - 'a'] == nil {
            return false
        }
        root = root.children[word[i] - 'a']
    }
    // 最后判断是前缀，还是完整单词
    if root.isEnd != true {
        return false
    }
    return true
}


/** Returns if there is any word in the trie that starts with the given prefix. */
func (this *Trie) StartsWith(prefix string) bool {
    root := this
    // 只判断前缀，不需要判断完整单词
    for i := 0; i < len(prefix); i++ {
        if root.children[prefix[i] - 'a'] == nil {
            return false
        }
        root = root.children[prefix[i] - 'a']
    }
    return true
}


/**
 * Your Trie object will be instantiated and called as such:
 * obj := Constructor();
 * obj.Insert(word);
 * param_2 := obj.Search(word);
 * param_3 := obj.StartsWith(prefix);
 */

类似二叉树的层序遍历

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
// dfs
func levelOrder(root *TreeNode) [][]int {
    res := [][]int{}
    if root == nil {
        return res
    }
    var dfs func(*TreeNode, int)
    dfs = func(root *TreeNode, l int) {
        if l == len(res) {
            res = append(res, []int{})
        }
        res[l] = append(res[l], root.Val)
        if root.Left != nil { dfs(root.Left, l + 1) }
        if root.Right != nil { dfs(root.Right, l + 1) }
    }
    dfs(root, 0)
    return res
}

// bfs
func levelOrder(root *TreeNode) [][]int {
    res := [][]int{}
    if root == nil {
        return res
    }
    q := []*TreeNode{root}

    for len(q) > 0 {
        qLen := len(q)
        curRes := []int{}
        for i := 0; i < qLen; i++ {
            node := q[0]
            q = q[1:]
            curRes = append(curRes, node.Val)
            if node.Left != nil { q = append(q, node.Left) }
            if node.Right != nil { q = append(q, node.Right) }
        }
        res = append(res, curRes)
    }
    return res
}

注意判断边界条件，开始的root是否为nil， root的left，right是否为nil

单词搜索

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
// dfs
func exist(board [][]byte, word string) bool {
    dirI := []int{-1, 1, 0, 0}
    dirJ := []int{0, 0, 1, -1}
    var dfs func(int, int, int) bool
    dfs = func(i, j, l int) bool {
        if board[i][j] != word[l] {
            return false
        }
        if l == len(word) - 1 {
            return true
        }
        board[i][j] = ' '
        defer func(i, j int, b byte)  { board[i][j] = b }(i, j, word[l])
        for k := 0; k < 4; k++ {
            newI := i + dirI[k]
            newJ := j + dirJ[k]
            if newI >= 0 && newJ >= 0 && newI < len(board) && newJ < len(board[0]) && board[newI][newJ] != ' ' {
                if dfs(newI, newJ, l + 1) { return true }
            }
        }
        return false
    }

    for i := 0; i < len(board); i++ {
        for j := 0; j < len(board[0]); j++ {
            if board[i][j] == word[0] {
                if dfs(i, j, 0) { return true }
            }
        }
    }
    return false
}

// Trie 字典树， 注释看单词搜索2，一样
type Trie struct {
    isEnd bool
    childrens [58]*Trie
}

func exist(board [][]byte, word string) bool {
    root := &Trie{}
    node := root
    for _, v := range word {
        if node.childrens[v - 'A'] == nil {
            node.childrens[v - 'A'] = &Trie{}
        }
        node = node.childrens[v - 'A']
    }
    node.isEnd = true

    dirx := []int{-1, 0, 1, 0}
    diry := []int{0, 1, 0, -1}

    var dfs func(x, y int, root *Trie) bool
    dfs = func(x, y int, root *Trie) bool {
        if x < 0 || y < 0 || x == len(board) || y == len(board[0]) {
            return false
        }
        cur := board[x][y]
        if cur == ' ' || root.childrens[cur - 'A'] == nil {
            return false
        }
        root = root.childrens[cur - 'A']
        if root.isEnd {
            return true
        }
        board[x][y] = ' '
        for p := 0; p < 4; p++ {
            if dfs(x + dirx[p], y + diry[p], root) {
                return true
            }
        }
        board[x][y] = cur
        return false
    }
    for x := 0; x < len(board); x++ {
        for y := 0; y < len(board[0]); y++ {
            if dfs(x, y, root) {
                return true
            }
        }
    }
    return false
}

单词搜索II

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
//dfs
func findWords(board [][]byte, words []string) []string {
    m := len(board)
    n := len(board[0])
    res := make([]string, 0, len(words))
    for i := 0; i < len(words); i++ {
        flag := false
        for j := 0; j < m; j++ {
            for k := 0; k < n; k++ {
                flag = dfs(board, words, j, k, i, 0)
                if flag { break }
            }
            if flag { break }
        }
        if flag { res = append(res, words[i]) }
    }
    return res
}

var dirx = []int{-1, 1, 0, 0}
var diry = []int{0, 0 ,-1, 1}

func dfs(board [][]byte, words []string, j, k, i, t int) bool {
    if board[j][k] != words[i][t] {
        return false
    }
    if t == len(words[i]) - 1 {
        return true
    }
    board[j][k] = ' '
    // 还原board
    defer func(j, k int, b byte) { board[j][k] = b } (j, k, words[i][t])
    for d := 0; d < 4; d++ {
        newJ := j + dirx[d]
        newK := k + diry[d]
        if newJ >= 0 && newK >= 0 && newJ < len(board) && newK < len(board[0]) && board[newJ][newK] != ' ' {
            if dfs(board, words, newJ, newK, i, t + 1) { return true }
        }
    }
    return false
}

// trie树方法
// 创建字典树
type Trie struct {
    word string
    childrens [26]*Trie
}
// 方向
var dirx = []int{-1, 1, 0, 0}
var diry = []int{0, 0, -1, 1}
// 主函数
func findWords(board [][]byte, words []string) []string {
    res := []string{}
    // 初始化root
    root := &Trie{}
    for _, v := range words {
    	// 每个单词重新从root开始
        node := root
        // 插入字典树
        for i := 0; i < len(v); i++ {
            if node.childrens[v[i] - 'a'] == nil {
                node.childrens[v[i] - 'a'] = &Trie{}
            }
            node = node.childrens[v[i] - 'a']
        }
        node.word = v
    }
    // 递归
    var dfs func(x, y int, root *Trie)
    dfs = func(x, y int, root *Trie) {
    	// 超出边界，return
        if x < 0 || y < 0 || x == len(board) || y == len(board[0]) {
            return
        }
        // 取当前点
        cur := board[x][y]
        // 如果当前点使用过，或者不在字典树里，则return
        if cur == ' ' || root.childrens[cur - 'a'] == nil {
            return
        }
        // 在字典树里，且没有使用过，则更新到下一个节点
        root = root.childrens[cur - 'a']
        // 如果到结束，找到单词，加入结果
        if root.word != "" {
            res = append(res, root.word)
            root.word = ""
        }
        // 标记防止重复
        board[x][y] = ' '
        // 继续dfs
        for p := 0; p < 4; p++ {
            dfs(x + dirx[p], y + diry[p], root)
        }
        // 还原
        board[x][y] = cur
    }
    // 遍历 dfs
    for x := 0; x < len(board); x++ {
        for y := 0; y < len(board[0]); y++ {
            dfs(x, y, root)
        }
    }
    return res
}

单词搜索II，字典树方法的时间复杂度
O（m * n * 4^L）m,n是二维网格的长宽，L是单词的平均长度