北川广海の梦

北川广海の梦

前缀树实现

2023-02-12

基于Golang,效率可能并不能达到O(N),因为要支持字母外的其他字符,所以子节点的长度并不固定,这里用的是链表存储。所以在进入到下一子节点前,需要遍历这个子节点链表。然而,如果是纯英文,则可以通过ASCII直接定位了,效率和内存都更有优势。

下面是代码:

package main

import (
	"container/list"
	"fmt"
)

func main() {
	var word = "hello world!"
	var root = NewTrie()
	root.Insert(word)

	result := root.Search("hello")
	fmt.Println(result)

	result = root.Search("world")
	fmt.Println(result)

}

type Node struct {
	value rune

	subs *list.List
}

func NewTrie() (root *Node) {
	return &Node{
		subs: list.New(),
	}
}
func makeNode(letter rune) *Node {
	return &Node{
		value: letter,
		subs:  list.New(),
	}
}
func (n *Node) Insert(word string) {
	cur := n
	for _, chr := range word {
		var matchNode *Node
		sub := cur.subs.Front()
		for {
			if sub == nil {
				break
			}
			// letter exists.
			node := sub.Value.(*Node)
			if node.value == chr {
				matchNode = node
				break
			}
			sub = sub.Next()
		}

		if matchNode == nil {
			// add letter into root
			matchNode = makeNode(chr)
			cur.subs.PushBack(matchNode)
		}
		cur = matchNode
	}

}

func (n *Node) Search(key string) bool {
	cur := n
	for _, chr := range key {
		var next *Node
		sub := cur.subs.Front()
		for {
			if sub == nil {
				break
			}
			// letter exists.
			node := sub.Value.(*Node)
			if node.value == chr {
				next = node
				break
			}
			sub = sub.Next()
		}
		if next == nil {
			return false
		}
		// go to next
		cur = next
	}
	return true
}