curl2info/structure.go

package curl2info

// TrieWord Trie 需要的Word接口
type TrieWord interface {
	GetWord() string
}

// TrieStrWord 最简单的TrieWord 结构
type TrieStrWord string

// GetWord 获取单词
func (tsw *TrieStrWord) GetWord() string {
	return (string)(*tsw)
}

// Trie 前缀树
type Trie struct {
	isWord bool
	value  interface{}
	char   byte
	prev   *Trie
	next   map[byte]*Trie
}

// NewTrie Initialize your data structure here.
func NewTrie() *Trie {
	return &Trie{next: make(map[byte]*Trie)}
}

// Insert a word into the trie.
func (trie *Trie) Insert(iword TrieWord) {
	cur := trie
	word := iword.GetWord()
	l := len(word)

	for i := 0; i < l; i++ {
		c := word[i]
		if next, ok := cur.next[c]; ok {
			cur = next
		} else {
			create := NewTrie()
			cur.next[c] = create
			create.char = c
			create.prev = cur
			cur = create
		}
	}

	cur.isWord = true
	cur.value = iword
}

// AllWords 所有单词
func (trie *Trie) AllWords() []string {
	var result []string
	for _, v := range trie.next {
		look(v, "", &result)
	}
	return result
}

func look(cur *Trie, content string, result *[]string) {
	content += string(cur.char)
	if cur.isWord {
		*result = append(*result, content)
	}
	for _, v := range cur.next {
		look(v, content, result)
	}
}

// Remove 移除单词
func (trie *Trie) Remove(word string) {
	cur := trie
	l := len(word)
	for i := 0; i < l; i++ {
		c := word[i]
		if next, ok := cur.next[c]; ok {
			cur = next
		} else {
			return
		}
	}

	if cur != nil {
		cur.isWord = false
		cur.value = nil

		lastchar := cur.char

		if len(cur.next) == 0 {
			for cur.isWord != true && cur.prev != nil {
				lastchar = cur.char
				cur = cur.prev
				if len(cur.next) > 1 {
					return
				}
			}
			delete(cur.next, lastchar)
		}
	}
}

// SearchMostPrefix Returns if the word is in the trie.
func (trie *Trie) SearchMostPrefix(iword TrieWord) interface{} {
	cur := trie
	word := iword.GetWord()

	l := len(word)

	var result interface{}
	for i := 0; i < l; i++ {
		c := word[i]
		if next, ok := cur.next[c]; ok {
			cur = next
			if cur.isWord {
				result = cur.value
			}
		} else {
			return result
		}
	}
	return result
}

// Match Returns if the word is in the trie.
func (trie *Trie) Match(iword TrieWord) interface{} {
	cur := trie
	word := iword.GetWord()

	l := len(word)
	for i := 0; i < l; i++ {
		c := word[i]
		if next, ok := cur.next[c]; ok {
			cur = next
		} else {
			return nil
		}
	}

	return cur.value
}

// StartsWith Returns if there is any word in the trie that starts with the given prefix. */
func (trie *Trie) StartsWith(prefix string) bool {
	cur := trie
	l := len(prefix)
	for i := 0; i < l; i++ {
		c := prefix[i]
		if next, ok := cur.next[c]; ok {
			cur = next
		} else {
			return false
		}
	}
	return true
}