[agent] - bastion mode init

This commit is contained in:
zeaslity
2024-04-17 15:02:41 +08:00
parent d54ebfc7b0
commit c0712b7325
3 changed files with 160 additions and 0 deletions

View File

@@ -0,0 +1,140 @@
package bastion_init
import (
"fmt"
"sort"
)
type TrieNode struct {
children [26]*TrieNode
isEnd bool
word string
}
type Trie struct {
root *TrieNode
}
func NewTrie() *Trie {
return &Trie{root: &TrieNode{}}
}
func (tn *TrieNode) Insert(word string) {
node := tn
for _, c := range word {
c -= 'a' // Convert char to int index
if node.children[c] == nil {
node.children[c] = &TrieNode{}
}
node = node.children[c]
}
node.isEnd = true
node.word = word
}
func (t *Trie) InsertAll(words []string) {
for _, word := range words {
t.root.Insert(word)
}
}
func (tn *TrieNode) Find(prefix string) []string {
node := tn
result := make([]string, 0)
for _, c := range prefix {
c -= 'a' // Convert char to int index
if node.children[c] != nil {
node = node.children[c]
} else {
return result // No more matching nodes
}
}
trieWalk(node, &result)
return result
}
func trieWalk(node *TrieNode, result *[]string) {
if node.isEnd {
*result = append(*result, node.word)
}
for _, child := range node.children {
if child != nil {
trieWalk(child, result)
}
}
}
type Word struct {
Word string
Rank int
}
func SortWords(words []string) []Word {
sortedWords := make([]Word, len(words))
for i, word := range words {
sortedWords[i] = Word{Word: word, Rank: i}
}
sort.Slice(sortedWords, func(i, j int) bool {
return sortedWords[i].Word < sortedWords[j].Word
})
return sortedWords
}
func FindClosestWord(trie *Trie, prefix string) (string, error) {
words := trie.root.Find(prefix)
if len(words) == 0 {
return "", fmt.Errorf("no words found for prefix: %s", prefix)
}
sortedWords := SortWords(words)
minDistance := len(prefix) // Initialize with the maximum possible distance
closestWord := sortedWords[0].Word
for _, word := range sortedWords {
distance := levDist(prefix, word.Word)
if distance < minDistance {
minDistance = distance
closestWord = word.Word
}
}
return closestWord, nil
}
func levDist(s1, s2 string) int {
m, n := len(s1), len(s2)
if m == 0 {
return n
}
if n == 0 {
return m
}
dp := make([][]int, m+1)
for i := range dp {
dp[i] = make([]int, n+1)
}
for i := 1; i <= m; i++ {
for j := 1; j <= n; j++ {
cost := 0
if s1[i-1] != s2[j-1] {
cost = 1
}
dp[i][j] = min(dp[i-1][j], dp[i][j-1], dp[i-1][j-1]) + cost
}
}
return dp[m][n]
}
func min(a, b, c int) int {
if a < b && a < c {
return a
} else if b < a && b < c {
return b
} else {
return c
}
}

View File

@@ -23,4 +23,5 @@ func PrintBastionHelp() {
} else {
fmt.Println("获取标准输出终端长度时发生错误:", err)
}
}