var myTree dt.Interface BeforeEach(func() { id3InputSet = il.InputLog{ {"1", "3"}, {"0", "1", "3"}, {"1", "3"}, {"0", "1", "3"}, {"1", "2"}, {"0", "1", "2"}, {}, {"0"}, } myTree = dt.New() myTree.Train(il.Interface(id3InputSet), "3") }) Describe("Training a Decision Tree", func() { It("it uses ID3 to maximize entropy", func() { Expect(myTree.Value()).To(Equal("2")) Expect(myTree.Left().Value()).To(Equal("1")) }) It("doesn't continue building a branch if a decision can be made", func() { Expect(myTree.Left().Right().Right()).To(BeNil()) }) It("builds a tree with the right decisions", func() { Expect(myTree.Left().Right().Value()).To(Equal(true)) Expect(myTree.Left().Left().Value()).To(Equal(false))
inputSet = il.InputLog{ {}, {"0", "1", "3"}, {"1", "3"}, {"0", "3"}, {"0", "1", "2", "3"}, {"0", "1", "3"}, {"0", "1"}, {"0", "2", "4"}, {"2", "4"}, {"4", "1"}, {"0", "4", "2"}, } myForest = pf.New() myForest.Train(il.Interface(inputSet)) }) Describe("Making a prediction", func() { It("adds a key that should be there", func() { Expect(myForest.Predict([]string{"0", "1"})).To(Equal([]string{"0", "1", "3"})) }) It("removes a key that doesn't belong", func() { Expect(myForest.Predict([]string{"0", "1", "2", "3"})).To(Equal([]string{"0", "1", "3"})) }) It("finds keys associated with a key", func() { Expect(myForest.Predict([]string{"4"})).To(Equal([]string{"0", "1", "2"})) }) })