Merge pull request #24 from IJNKAWAKAZE/dev

使用后缀树来提高inline模式干员和敌人的查询效率 O(n^2) -> O(n)

src/plugins/datasource/update_data_source.go

@@ -201,11 +201,11 @@ func UpdateDataSourceRunner() {
 	}
 
 	defer response.Body.Close()
-	utils.OperatorMap = make(map[string]utils.Operator)
-	utils.RecruitOperatorList = nil
+
 	utils.RedisSet("operatorList", json.MustMarshalString(operators), 0)
 	MaterialInfo()
 	log.Println("数据源更新完毕")
+	utils.DataNeedUpdate = true
 }
 
 func MaterialInfo() {

src/utils/arknights_utils.go

@@ -1,11 +1,13 @@
 package utils
 
 import (
+	"arknights_bot/utils/suffixtree"
 	"encoding/json"
 	"fmt"
 	"github.com/spf13/viper"
 	"github.com/tidwall/gjson"
 	"io"
+	"log"
 	"net/http"
 	"net/url"
 	"strings"
@@ -55,64 +57,111 @@ type Material struct {
 	StageEfficiency   string `json:"stageEfficiency"`   // 关卡效率
 }
 
-var OperatorMap = make(map[string]Operator)
-var RecruitOperatorList []Operator
+var operatorMap = make(map[string]Operator)
+var recruitOperatorList []Operator
+var DataNeedUpdate = true
+var operatorTree suffixtree.GST
+var itemTree suffixtree.GST
+var itemArray []string
+var enemyTree suffixtree.GST
+var enemyArray []pair
+var operators []Operator
 
 func GetOperators() []Operator {
-	var operators []Operator
+	updateData()
+	return operators
+}
+func updateData() {
+	if !DataNeedUpdate {
+		return
+	}
+	//operators
 	operatorsJson := RedisGet("operatorList")
 	json.Unmarshal([]byte(operatorsJson), &operators)
-	return operators
+	operatorMap = make(map[string]Operator)
+	operatorTree = suffixtree.NewGeneralizedSuffixTree()
+	for index, operator := range operators {
+		operatorMap[strings.ToLower(operator.Name)] = operator
+		operatorTree.Put(strings.ToLower(operator.Name), index)
+		if strings.Contains(operator.ObtainMethod, "公开招募") {
+			recruitOperatorList = append(recruitOperatorList, operator)
+		}
+	}
+	//enemy
+	func() {
+		resultArray, resultTree := fetchEnemiesData()
+		enemyArray = resultArray
+		enemyTree = resultTree
+		defer func() {
+			if err := recover(); err != nil {
+				log.Fatal("Can not update enemy")
+			}
+		}()
+	}()
+
+	//items
+	//gjson.Parse(RedisGet("materialMap"))
+	//NOT_TODO: check what is the materialMap 暂时不做因为yj还没搞出太多材料
+	// 当前复杂度还能用
+	//set flag
+	DataNeedUpdate = false
+
 }
 
-func GetOperatorByName(name string) Operator {
-	if len(OperatorMap) == 0 {
-		for _, op := range GetOperators() {
-			OperatorMap[op.Name] = op
-		}
+type pair struct {
+	a, b interface{}
+}
+
+func fetchEnemiesData() ([]pair, suffixtree.GST) {
+	makeurl := func(n string) string {
+		paintingName := fmt.Sprintf("头像_敌人_%s.png", n)
+		m := Md5(paintingName)
+		path := "https://media.prts.wiki" + fmt.Sprintf("/%s/%s/", m[:1], m[:2])
+		return path + url.PathEscape(paintingName)
 	}
-	return OperatorMap[name]
+	emeryTree := suffixtree.NewGeneralizedSuffixTree()
+	var newEnemyArray []pair
+	api := viper.GetString("api.enemy")
+	response, _ := http.Get(api)
+	e, _ := io.ReadAll(response.Body)
+	defer response.Body.Close()
+	enemyJson := gjson.ParseBytes(e)
+	for index, en := range enemyJson.Array() {
+		n := en.Get("name").String()
+		newEnemyArray = append(newEnemyArray, pair{n, makeurl(n)})
+		emeryTree.Put(strings.ToLower(n), index)
+	}
+	return newEnemyArray, emeryTree
+}
+func GetOperatorByName(name string) Operator {
+	updateData()
+	return operatorMap[name]
 }
 
 func GetOperatorsByName(name string) []Operator {
+	updateData()
 	var operatorList []Operator
-	for _, op := range GetOperators() {
-		if strings.Contains(strings.ToLower(op.Name), strings.ToLower(name)) {
-			operatorList = append(operatorList, op)
-		}
+	for _, op := range operatorTree.Search(strings.ToLower(name)) {
+		print(operators[op].Name)
+		operatorList = append(operatorList, operators[op])
 	}
+	println()
 	return operatorList
 }
 
 func GetRecruitOperatorList() []Operator {
-	if len(RecruitOperatorList) == 0 {
-		for _, op := range GetOperators() {
-			if strings.Contains(op.ObtainMethod, "公开招募") {
-				RecruitOperatorList = append(RecruitOperatorList, op)
-			}
-		}
-	}
-	return RecruitOperatorList
+	updateData()
+	return recruitOperatorList
 }
 
 func GetEnemiesByName(name string) map[string]string {
-	var enemyList = make(map[string]string)
-	api := viper.GetString("api.enemy")
-	response, _ := http.Get(api)
-	e, _ := io.ReadAll(response.Body)
-	defer response.Body.Close()
-	enemyJson := gjson.ParseBytes(e)
-	for _, en := range enemyJson.Array() {
-		n := en.Get("name").String()
-		if strings.Contains(strings.ToLower(n), strings.ToLower(name)) {
-			paintingName := fmt.Sprintf("头像_敌人_%s.png", n)
-			m := Md5(paintingName)
-			path := "https://media.prts.wiki" + fmt.Sprintf("/%s/%s/", m[:1], m[:2])
-			pic := path + url.PathEscape(paintingName)
-			enemyList[n] = pic
-		}
+	updateData()
+	var enemyMap = make(map[string]string)
+	for _, index := range enemyTree.Search(strings.ToLower(name)) {
+		a := enemyArray[index]
+		enemyMap[a.a.(string)] = a.b.(string)
 	}
-	return enemyList
+	return enemyMap
 }
 
 func GetItemsByName(name string) map[string]string {

src/utils/suffixtree/README.txt

@@ -0,0 +1 @@
+from https://github.com/ljfuyuan/suffixtree modified to make the type is usable

src/utils/suffixtree/edge.go

@@ -0,0 +1,10 @@
+package suffixtree
+
+type edge struct {
+	label []rune
+	*node
+}
+
+func newEdge(label []rune, node *node) *edge {
+	return &edge{label: label, node: node}
+}

src/utils/suffixtree/node.go

@@ -0,0 +1,102 @@
+package suffixtree
+
+import (
+	"sort"
+)
+
+type node struct {
+	/*
+	 * The payload array used to store the data (indexes) associated with this node.
+	 * In this case, it is used to store all property indexes.
+	 */
+	data []int
+	/**
+	 * The set of edges starting from this node
+	 */
+	edges []*edge
+	/**
+	 * The suffix link as described in Ukkonen's paper.
+	 * if str is the string denoted by the path from the root to this, this.suffix
+	 * is the node denoted by the path that corresponds to str without the first rune.
+	 */
+	suffix *node
+}
+
+/*
+ * getData returns the first numElements elements from the ones associated to this node.
+ *
+ * Gets data from the payload of both this node and its children, the string representation
+ * of the path to this node is a substring of the one of the children nodes.
+ *
+ * @param numElements the number of results to return. Use <=0 to get all
+ * @return the first numElements associated to this node and children
+ */
+func (n *node) getData() (ret []int) {
+
+	ret = n.data
+
+	// need to get more matches from child nodes. This is what may waste time
+	for _, edge := range n.edges {
+		data := edge.node.getData()
+		ret = append(ret, data...)
+
+	}
+	return
+}
+
+// addRef adds the given index to the set of indexes associated with this
+func (n *node) addRef(index int) {
+	if n.contains(index) {
+		return
+	}
+	n.addIndex(index)
+	// add this reference to all the suffixes as well
+	iter := n.suffix
+	for iter != nil {
+		if iter.contains(index) {
+			break
+		}
+		iter.addRef(index)
+		iter = iter.suffix
+	}
+}
+
+func (n *node) contains(index int) bool {
+	i := sort.SearchInts(n.data, index)
+	return i < len(n.data) && n.data[i] == index
+}
+
+func (n *node) addEdge(r rune, e *edge) {
+	if idx := n.search(r); idx == -1 {
+		n.edges = append(n.edges, e)
+		sort.Slice(n.edges, func(i, j int) bool { return n.edges[i].label[0] < n.edges[j].label[0] })
+	} else {
+		n.edges[idx] = e
+	}
+
+}
+
+func (n *node) getEdge(r rune) *edge {
+	idx := n.search(r)
+	if idx < 0 {
+		return nil
+	}
+	return n.edges[idx]
+}
+
+func (n *node) search(r rune) int {
+	idx := sort.Search(len(n.edges), func(i int) bool { return n.edges[i].label[0] >= r })
+	if idx < len(n.edges) && n.edges[idx].label[0] == r {
+		return idx
+	}
+
+	return -1
+}
+
+func (n *node) addIndex(idx int) {
+	n.data = append(n.data, idx)
+}
+
+func newNode() *node {
+	return &node{}
+}