spider_creater/backend/prefixtree.go

package backend

import (
	"strings"
	"unicode/utf8"
)

// 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,
	}
}

// Trie represents the Trie data structure
type Trie struct {
	root *TrieNode
}

// NewTrie creates a new Trie
func NewTrie() *Trie {
	return &Trie{root: NewTrieNode()}
}

// Insert inserts a word into the Trie
func (t *Trie) Insert(word string) {
	node := t.root
	for _, char := range word {
		if _, ok := node.children[char]; !ok {
			node.children[char] = NewTrieNode()
		}
		node = node.children[char]
	}
	node.isEnd = true
}

// SearchInText searches for words in the text that are present in the Trie
func (t *Trie) SearchInText(text string) []string {
	var foundWords []string
	node := t.root
	var wordBuilder []rune

	for i, w := 0, 0; i < len(text); i += w {
		runeValue, width := utf8.DecodeRuneInString(text[i:])
		w = width

		if runeValue == utf8.RuneError {
			continue
		}

		if _, ok := node.children[runeValue]; !ok {
			node = t.root
			wordBuilder = wordBuilder[:0]
			continue
		}

		node = node.children[runeValue]
		wordBuilder = append(wordBuilder, runeValue)

		if node.isEnd {
			foundWords = append(foundWords, string(wordBuilder))
			// Continue to search for other words starting from this node
			node = t.root
			wordBuilder = wordBuilder[:0]
		}
	}

	return foundWords
}

// SearchInText searches for words in the text that are present in the Trie
func (t *Trie) HasKeyword(text string) bool {
	node := t.root
	var wordBuilder []rune

	for i, w := 0, 0; i < len(text); i += w {
		runeValue, width := utf8.DecodeRuneInString(text[i:])
		w = width
		if runeValue == utf8.RuneError {
			continue
		}

		if _, ok := node.children[runeValue]; !ok {
			node = t.root
			wordBuilder = wordBuilder[:0]
			continue
		}

		node = node.children[runeValue]
		wordBuilder = append(wordBuilder, runeValue)

		if node.isEnd {
			return true
		}
	}

	return false
}

// BatchInsert
func (t *Trie) BatchInsert(words string) {
	for _, s := range strings.Split(words, ";") {
		if s != "" {
			t.Insert(s)
		}
	}
}