func TrainNB0(trainMatrix [][]int, trainCategory []int) ([]float64, []float64, float64) {
numTrainDocs := len(trainMatrix)
numWords := len(trainMatrix[0])
sumCategory := calc.SumInt(trainCategory)
pAbusive := float64(sumCategory) / float64(numTrainDocs)
p0Num := make([]int, numWords)
p1Num := make([]int, numWords)
p0Denom := float64(0.0)
p1Denom := float64(0.0)
for i := 0; i < numTrainDocs; i++ {
if trainCategory[i] == 1 {
p1Num = calc.AddIntVec(p1Num, trainMatrix[i])
p1Denom += float64(calc.SumInt(trainMatrix[i]))
} else {
p0Num = calc.AddIntVec(p0Num, trainMatrix[i])
p0Denom += float64(calc.SumInt(trainMatrix[i]))
}
}
fmt.Println("p0Num: ", p0Num)
fmt.Println("p1Num: ", p1Num)
p1Vec := calc.DivIntVec(p1Num, p1Denom)
p0Vec := calc.DivIntVec(p0Num, p0Denom)
return p0Vec, p1Vec, pAbusive
}
func Test_AddIntVec(t *testing.T) {
a := []int{1, 2, 3}
b := []int{4, 5, 6}
result := calc.AddIntVec(a, b)
if len(result) != len(a) {
t.Fatalf("Add wrong")
}
for i, count := 0, len(a); i < count; i++ {
if result[i] != a[i]+b[i] {
t.Fatalf("Add to wrong result")
}
}
}