jyseo/utility/match.go

package utility

import (
	"fmt"
	"strings"
)

// TrieNode represents each node in the Trie
type TrieNode struct {
	children map[rune]*TrieNode
	isEnd    bool
}

// NewTrieNode creates a new TrieNode
func NewTrieNode() *TrieNode {
	return &TrieNode{
		children: make(map[rune]*TrieNode),
		isEnd:    false,
	}
}

// 批量追加词
func (n *TrieNode) Insert(words ...string) {
	for _, word := range words {
		node := n
		for _, char := range word {
			if _, ok := node.children[char]; !ok {
				node.children[char] = NewTrieNode()
			}
			node = node.children[char]
		}
		node.isEnd = true
	}
}

// 前缀树替换所有
func (n *TrieNode) ReplaceAll(text, str string) string {
	var rs = []rune(text)
	var sb = new(strings.Builder)
	var lastPos = 0
	for i := 0; i < len(rs); i++ {
		node := n
		for j := i; j < len(rs); j++ {
			char := rs[j]
			if _, ok := node.children[char]; !ok {
				break
			}
			node = node.children[char]
			if node.isEnd {
				//前部分
				sb.WriteString(string(rs[lastPos:i]))
				sb.WriteString(fmt.Sprintf(" %s ", str))
				lastPos = j + 1
				i = j
				break
			}
		}
	}
	if lastPos == 0 { //修正全文未匹配情况
		return text
	} else if lastPos < len(rs) { //修正最后截取尾巴部分
		sb.WriteString(string(rs[lastPos:]))
	}
	return sb.String()
}