数据结构和算法学习笔记（五）前缀树和贪心算法

  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

package trie

import (
	"errors"
	"math"
)

type Node struct {
	Pass  int     // 添加元素时，node 走过几次
	End   int     // 有几个元素是以当前 node 最为结尾的
	Nexts []*Node // 下级 node 切片, 如果字符太多，可以用 map 代替
}

var (
	ErrNotContain = errors.New("not in trie")
)

func NewNode() *Node {
	return &Node{
		Pass:  0,
		End:   0,
		Nexts: make([]*Node, math.MaxUint8),
	}
}

type Tree struct {
	Root *Node // Root 节点表示空字符串，如过 Root.End == 1，那表示加入了一个空字符串
}

func NewTree() *Tree {
	return &Tree{Root: NewNode()}
}

// Insert 将单词插入前缀树
func (t *Tree) Insert(word string) error {
	chars := []byte(word)

	node := t.Root
	node.Pass++
	for _, c := range chars {
		index := int(c) // 0 位置表示 'a', 1 位置表示 'b'...
		if node.Nexts[index] == nil {
			node.Nexts[index] = NewNode() // 没有就新建
		}
		node = node.Nexts[index]
		node.Pass++ // 路过这个 node, Pass++
	}
	node.End++ // word 放入完毕，最后一个 node 的 End++
	return nil
}

// lastNode 返回单词在树中的最后一个 node
func (t *Tree) lastNode(word string) (*Node, error) {
	chars := []byte(word)

	node := t.Root
	for _, c := range chars {
		index := int(c)
		if node.Nexts[index] == nil {
			return nil, ErrNotContain
		}
		node = node.Nexts[index]
	}
	return node, nil
}

// Search 返回单词在前缀树中插入了几次
func (t *Tree) Search(word string) int {
	lastNode, err := t.lastNode(word)
	if err != nil {
		return 0
	}
	return lastNode.End
}

// PrefixNumber 返回有几个以 pre 作前缀的字符串，注意一个字符串可能加入多次
func (t *Tree) PrefixNumber(pre string) int {
	lastNode, err := t.lastNode(pre)
	if err != nil {
		return 0
	}
	return lastNode.Pass
}

// Delete 从前缀树中删除单词
func (t *Tree) Delete(word string) {
	_, err := t.lastNode(word)
	if err != nil {
		return
	}
	chars := []byte(word)

	node := t.Root
	node.Pass--
	for _, c := range chars {
		index := int(c)
		if node.Nexts[index].Pass--; node.Nexts[index].Pass == 0 {
			node.Nexts[index] = nil
			return
		}
		node = node.Nexts[index]
	}
	node.End--
}

  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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

package trie_test

import (
	"math/rand"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/boh5/dsal/trie"
)

const alphabet = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

func init() {
	rand.Seed(int64(time.Now().Nanosecond()))
}
func insertTree(t *testing.T, tree *trie.Tree, words []string, numInTree int) {
	require.NotNil(t, tree)
	for i, w := range words {
		err := tree.Insert(w)
		require.NoError(t, err)
		require.Equal(t, numInTree+i+1, tree.Root.Pass)
	}
	require.Equal(t, 0, tree.Root.End)
}

func TestTree_Insert(t *testing.T) {
	t.Parallel()

	n := 10
	words := randomUniqueWords(n, 10)
	tree := trie.NewTree()

	insertTree(t, tree, words, 0)
	insertTree(t, tree, words, n)
}

func TestTree_Search(t *testing.T) {
	t.Parallel()

	n := 10
	words := randomUniqueWords(n, 10)
	tree := trie.NewTree()

	insertTree(t, tree, words, 0)

	for _, word := range words {
		i := tree.Search(word)
		require.Equal(t, 1, i)
	}

	wordsNotInTree := []string{"NOT", "IN", "Tree"}
	for _, word := range wordsNotInTree {
		i := tree.Search(word)
		require.Equal(t, 0, i)
	}

	repeatWords := words[:5]
	insertTree(t, tree, repeatWords, n)
	for _, word := range repeatWords {
		i := tree.Search(word)
		require.Equal(t, 2, i)
	}
}

func TestTree_PrefixNumber(t *testing.T) {
	t.Parallel()

	words := []string{"abc", "abcd", "abcde", "abcdef"}
	n := len(words)
	tree := trie.NewTree()

	insertTree(t, tree, words, 0)

	for i, word := range words {
		ans := tree.PrefixNumber(word)
		require.Equal(t, n-i, ans)
	}

	notContainPrefixes := []string{"123", "eee", "xyz"}
	for _, word := range notContainPrefixes {
		ans := tree.PrefixNumber(word)
		require.Equal(t, 0, ans)
	}
}

func TestTree_Delete(t *testing.T) {
	t.Parallel()

	n := 10
	words := randomUniqueWords(n, 10)
	tree := trie.NewTree()

	insertTree(t, tree, words, 0)

	// Delete not exists
	tree.Delete("NE")
	for _, word := range words {
		i := tree.Search(word)
		require.Equal(t, 1, i)
	}

	// Delete repeat
	insertTree(t, tree, words[:1], n)
	tree.Delete(words[0])
	require.Equal(t, 1, tree.Search(words[0]))

	// Delete
	for _, word := range words {
		tree.Delete(word)
		i := tree.Search(word)
		require.Equal(t, 0, i)
	}

	// Last status
	require.Equal(t, 0, tree.Root.Pass)
	for _, node := range tree.Root.Nexts {
		require.Nil(t, node)
	}

}

func randomString(n int) string {
	var sb strings.Builder
	for i := 0; i < n; i++ {
		sb.WriteByte(alphabet[rand.Intn(len(alphabet))])
	}
	return sb.String()
}

func randomUniqueWords(n, wordLen int) []string {
	words := make([]string, n)
	seen := make(map[string]any)
	for i := 0; i < n; i++ {
		var str string
		for {
			str = randomString(wordLen)
			if _, ok := seen[str]; !ok {
				break
			}
		}
		seen[str] = nil
		words[i] = str
	}
	return words
}