golang-实现一个泛型版本的前缀树trie

Trie

当我们在做key-value存储时，最简单的就是 map[string]string, 假如有两个键值对('child', 'value1'), ('children', 'value2')，毫无疑问，value1与value2是一定要存的，不过 'child'与‘children'不一定都要存下来，只存'children'也不是不可以，这样就少存一个'child'，省点空间，不过我们要在字符‘d' 和 'n'处做上记号。

code

package trie

type Trie[T any] struct {
	isEnd    bool  // 只有当isEnd为true时，value字段才会存值
	value    T
	children map[rune]*Trie[T]
}

func New[T any]() *Trie[T] {
	return &Trie[T]{
		isEnd:    false,
		children: map[rune]*Trie[T]{},
	}
}

func (root *Trie[T]) Insert(word string, value T) {
	node := root
	for _, c := range word {
		if child, ok := node.children[c]; ok {
			node = child
		} else {
			newNode := &Trie[T]{
				children: map[rune]*Trie[T]{},
			}
			node.children[c] = newNode
			node = newNode
		}
	}

	node.isEnd = true
	node.value = value
}

func (root *Trie[T]) Search(word string) (isEnd bool, value T) {
	node := root
	for _, c := range word {
		if child, ok := node.children[c]; ok {
			node = child
			continue
		} else {
			isEnd = false
			return
		}
	}
	isEnd = node.isEnd
	value = node.value
	return
}

func (root *Trie[T]) StartWith(prefix string) bool {
	node := root
	for _, c := range prefix {
		if child, ok := node.children[c]; ok {
			node = child
			continue
		} else {
			return false
		}
	}
	return true
}

Test

package trie

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

// go test -run TestTrie
func TestTrie(t *testing.T) {
	trie := New[int]()

	word1, value1 := "letter", 25

	trie.Insert(word1, value1)

	ok1, ret1 := trie.Search("lette")

	assert.Equal(t, ok1, false)
	assert.Equal(t, ret1, 0)

	ok2 := trie.StartWith("lette")
	assert.Equal(t, ok2, true)

	ok3, ret2 := trie.Search("letter")

	assert.Equal(t, ok3, true)
	assert.Equal(t, ret2, 25)

	word2, value2 := "lette", 26

	trie.Insert(word2, value2)

	ok4, ret3 := trie.Search("lette")
	assert.Equal(t, ok4, true)
	assert.Equal(t, ret3, 26)

}

SourceCode

leetcode208

type Trie struct {
    isWord bool
    children map[rune]*Trie
}


func Constructor() Trie {
    return Trie{
        isWord: false,
        children: make(map[rune]*Trie),
    }
}


func (this *Trie) Insert(word string)  {
    node := this
    for _, c := range word {
        if child, ok := node.children[c]; ok {
            node = child
        } else {
            newNode := &Trie{
                isWord: false,
                children: make(map[rune]*Trie),
            }
            node.children[c] = newNode
            node = newNode
        }
    }
    node.isWord = true
}


func (this *Trie) Search(word string) bool {
    node := this
    for _, c := range word {
        if child, ok := node.children[c]; ok {
            node = child
            continue
        }
        return false
    }
    return node.isWord
}


func (this *Trie) StartsWith(prefix string) bool {
    node := this
    for _, c := range prefix {
        if child, ok := node.children[c]; ok {
            node = child
            continue
        }
        return false
    }
    return true
}