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"}))
})
})