2022-05-14：语法补全功能，比如"as soon as possible"，当我们识别到"as soon as"时, 基

package main

import (
  "fmt"
  "sort"
  "strings"
)

func main() {
  sentences := []string{"i think you are good", "i think you are fine", "i think you are good man"}
  k := 2
  ai := NewAI(sentences, k)
  for _, ans := range ai.suggest("i think you are") {
    fmt.Println(ans)
  }
  fmt.Println("=====")
  ai.fill("i think you are fucking good")
  ai.fill("i think you are fucking great")
  ai.fill("i think you are fucking genius")
  for _, ans := range ai.suggest("i think you are") {
    fmt.Println(ans)
  }
  fmt.Println("=====")
}

type TrieNode struct {
  word      string
  times     int
  nextNodes map[string]*TrieNode
  nextRanks []*TrieNode
}

func NewTrieNode(w string) *TrieNode {
  ans := &TrieNode{}
  ans.word = w
  ans.times = 1
  ans.nextNodes = make(map[string]*TrieNode)
  ans.nextRanks = make([]*TrieNode, 0)
  return ans
}

type AI struct {
  root *TrieNode
  topk int
}

func NewAI(sentences []string, k int) *AI {
  ans := &AI{}
  ans.root = NewTrieNode("")
  ans.topk = k
  for _, sentence := range sentences {
    ans.fill(sentence)
  }
  return ans
}

func (this *AI) fill(sentence string) {
  cur := this.root
  var next *TrieNode
  for _, word := range strings.Split(sentence, " ") {
    if _, ok := cur.nextNodes[word]; !ok {
      next = NewTrieNode(word)
      cur.nextNodes[word] = next
      cur.nextRanks = append(cur.nextRanks, next)
    } else {
      next = cur.nextNodes[word]
      next.times++
    }
    cur = next
  }
}

func (this *AI) suggest(sentence string) []string {
  ans := make([]string, 0)
  cur := this.root
  for _, word := range strings.Split(sentence, " ") {
    if _, ok := cur.nextNodes[word]; !ok {
      return ans
    } else {
      cur = cur.nextNodes[word]
    }
  }
  sort.Slice(cur.nextRanks, func(i, j int) bool {
    a := cur.nextRanks[i]
    b := cur.nextRanks[j]
    if a.times != b.times {
      return a.times > b.times
    } else {
      return a.word < b.word
    }
  })
  for _, n := range cur.nextRanks {
    ans = append(ans, n.word)
    if len(ans) == this.topk {
      break
    }
  }
  return ans
}