go获取tfidf词频矩阵

悠悠 2022-09-09 14:53 205阅读 0赞

package utils
    
    import (
    	"math"
    	"sort"
    	"time"
    )
    
    type TfidfUtil struct{ }
    
    type wordTfidf struct { 
    	word      string
    	frequency float64
    }
    
    type wordTfidfs []wordTfidf
    
    type Interface interface { 
    	Len() int
    	Less(i, j int) bool
    	Swap(i, j int)
    }
    
    func (wts wordTfidfs) Len() int { 
    	return len(wts)
    }
    func (wts wordTfidfs) Less(i, j int) bool { 
    	return wts[i].frequency < wts[j].frequency
    }
    func (wts wordTfidfs) Swap(i, j int) { 
    	wts[i], wts[j] = wts[j], wts[i]
    }
    
    func (wts *TfidfUtil) currentTimeMillis() int64 { 
    	return time.Now().UnixNano() / 1000000
    }
    
    func (wts *wordTfidfs) Sort() { 
    	sort.Sort(wts)
    }
    
    // listWords：分词后的矩阵
    // 获取词频矩阵 [[key, frequency]]
    func (tu *TfidfUtil) fit(listWords [][]string) wordTfidfs { 
    	docFrequency := make(map[string]float64, 0)
    	sumWorlds := 0
    	for _, wordList := range listWords { 
    		for _, v := range wordList { 
    			docFrequency[v] += 1
    			sumWorlds++
    		}
    	}
    	wordTf := make(map[string]float64)
    	for k, _ := range docFrequency { 
    		wordTf[k] = docFrequency[k] / float64(sumWorlds)
    	}
    	docNum := float64(len(listWords))
    	wordIdf := make(map[string]float64)
    	wordDoc := make(map[string]float64, 0)
    	for k, _ := range docFrequency { 
    		for _, v := range listWords { 
    			for _, vs := range v { 
    				if k == vs { 
    					wordDoc[k] += 1
    					break
    				}
    			}
    		}
    	}
    	for k, _ := range docFrequency { 
    		wordIdf[k] = math.Log(docNum / (wordDoc[k] + 1))
    	}
    	var words wordTfidfs
    	for k, _ := range docFrequency { 
    		words = append(words, wordTfidf{ 
    			word:      k,
    			frequency: wordTf[k] * wordIdf[k],
    		})
    	}
    	return words
    }

执行测试代码

func Load() [][]string { 
    	slice := [][]string{ 
    		{ "my", "dog", "has", "flea", "problems", "help", "please"},
    		{ "maybe", "not", "take", "him", "to", "dog", "park", "stupid"},
    		{ "my", "dalmation", "is", "so", "cute", "I", "love", "him"},
    		{ "stop", "posting", "stupid", "worthless", "garbage"},
    		{ "mr", "licks", "ate", "my", "steak", "how", "to", "stop", "him"},
    		{ "quit", "buying", "worthless", "dog", "food", "stupid"},
    	}
    	return slice
    }
    
    func TestTfIdt(t *testing.T) { 
    	var tf TfidfUtil
    	res := tf.fit(Load())
    	fmt.Println(res)
    }

![在这里插入图片描述][fa18c034b44b4400b37b973cb01867f9.png]

[fa18c034b44b4400b37b973cb01867f9.png]: /images/20220829/17d59c4056634beea25fed26be82f6c0.png