keyword-extraction/keywordprocessor.go

package flashtext

import (
	"unicode"
	"unicode/utf8"
)

type WalkFn func(start, end int) bool

type KeywordProcessor struct {
	root              *Node
	whiteSpaceChars   []string // 不知道有啥用，反正python也定义了，就跟着定义一个吧
	nonWordBoundaries string
	caseSensitive     bool // 匹配是否区分大小写
}

func NewKeywordProcessor(caseSensitive bool) *KeywordProcessor {
	return &KeywordProcessor{
		root:            newNode(),
		caseSensitive:   caseSensitive,
		whiteSpaceChars: []string{".", "\t", "\n", "\a", " ", ","}, // python 就是这样定义的，而且它也没有用到，所以先留着这里吧
	}
}

func (kp *KeywordProcessor) setItem(keyword string) {
	if len(keyword) == 0 {
		return
	}

	node := kp.root
	for _, char := range keyword {
		if !kp.caseSensitive {
			char = unicode.ToLower(char)
		}
		if _, ok := node.children[char]; !ok {
			// 子节点已存在 将子节点设为当前节点
			node.children[char] = newNode()
		}
		node = node.children[char]
	}
	// 记录当前匹配词的长度
	node.exist[len(keyword)] = struct{}{}
}

func (kp *KeywordProcessor) Build() {
	// 创建一个空队列，用于层级遍历 Trie 树节点
	queue := make([]*Node, 0)
	// 将根节点作为初始节点
	queue = append(queue, kp.root)

	for len(queue) > 0 {
		// 获取队列中的第一个节点
		currentNode := queue[0]
		// 弹出队列中的第一个节点
		queue = queue[1:]

		// 遍历当前节点的子节点
		for char, childNode := range currentNode.children {
			// 将子节点添加到队列中
			queue = append(queue, childNode)
			// 当前节点父节点的失败指针
			faFail := currentNode.failure

			// 递归获取父失败指针 直到获取到根节点为止
			for faFail != nil && faFail.children[char] == nil {
				faFail = faFail.failure
			}
			childNode.failure = kp.root
			if faFail != nil {
				childNode.failure = faFail.children[char]
			}
			for key := range childNode.failure.exist {
				childNode.exist[key] = struct{}{}
			}
		}
	}
}
func (kp *KeywordProcessor) AddKeyWord(keyword string) *KeywordProcessor {
	kp.setItem(keyword)
	return kp
}

func (kp *KeywordProcessor) AddKeywordsFromList(keywords []string) *KeywordProcessor {
	for _, keyword := range keywords {
		kp.setItem(keyword)
	}
	return kp
}

func (kp *KeywordProcessor) walk(sentence string, wf WalkFn) {
	// 从根节点开始查找
	currentNode := kp.root
	// 遍历文本 sentence 的每个字符，并记录当前字符的索引为 idx，当前字符为 r
	idx := 0
	for len(sentence) > 0 {
		r, l := utf8.DecodeRuneInString(sentence)
		idx += l
		sentence = sentence[l:]
		if !kp.caseSensitive {
			r = unicode.ToLower(r)
		}
		// 在循环中 判断是否有当前字符子节点
		//如果不存在，则说明匹配失败，需要通过失败路径回溯到前一个节点，直到找到一个匹配的子节点或回溯到根节点。
		for currentNode.children[r] == nil && currentNode.failure != nil {
			currentNode = currentNode.failure
		}
		if currentNode.children[r] == nil {
			continue
		}
		currentNode = currentNode.children[r]
		for length := range currentNode.exist {
			if !wf(idx-length, idx) {
				return
			}
		}
	}
}

func (kp *KeywordProcessor) walkByte(sentence []byte, wf WalkFn) {
	// 从根节点开始查找
	currentNode := kp.root
	// 遍历文本 sentence 的每个字符，并记录当前字符的索引为 idx，当前字符为 r
	idx := 0
	for len(sentence) > 0 {
		r, l := utf8.DecodeRune(sentence)
		idx += l
		sentence = sentence[l:]
		if !kp.caseSensitive {
			r = unicode.ToLower(r)
		}
		// 在循环中 判断是否有当前字符子节点
		//如果不存在，则说明匹配失败，需要通过失败路径回溯到前一个节点，直到找到一个匹配的子节点或回溯到根节点。
		for currentNode.children[r] == nil && currentNode.failure != nil {
			currentNode = currentNode.failure
		}
		if currentNode.children[r] == nil {
			continue
		}
		currentNode = currentNode.children[r]
		for length := range currentNode.exist {
			if !wf(idx-length, idx) {
				return
			}
		}
	}
}

// ExtractKeywords 匹配关键词
func (kp *KeywordProcessor) ExtractKeywords(sentence string) []Match {
	var matches []Match
	if len(sentence) == 0 {
		return matches
	}

	kp.walk(sentence, func(start, end int) bool {
		matches = append(matches, Match{
			start: start,
			end:   end,
			match: sentence[start:end],
		})
		return true
	})
	return matches
}

func (kp *KeywordProcessor) ExtractKeywordsFromBytes(sentence []byte) []Match {
	var matches []Match
	if len(sentence) == 0 {
		return matches
	}

	kp.walkByte(sentence, func(start, end int) bool {
		matches = append(matches, Match{
			start: start,
			end:   end,
			match: string(sentence[start:end]),
		})
		return true
	})
	return matches
}