def testSimpleCase(self): tree = DecisionTree() tree.fit(simpleData) for datum in simpleData: self.assertEqual(datum[-1], tree.predict(datum)) tree.print()
def testComplexCase(self): tree = DecisionTree() tree.fit(complexData) for datum in complexData: self.assertEqual(datum[-1], tree.predict(datum)) # Overcast = YES self.assertEqual( Label.YES, tree.predict( [Outlook.Overcast, Temperature.Hot, Humidity.High, Wind.Weak])) self.assertEqual( Label.YES, tree.predict([ Outlook.Overcast, Temperature.Cool, Humidity.Normal, Wind.Strong ])) # Sunny + Normal = YES self.assertEqual( Label.YES, tree.predict([ Outlook.Sunny, Temperature.Cool, Humidity.Normal, Wind.Strong ])) self.assertEqual( Label.YES, tree.predict( [Outlook.Sunny, Temperature.Hot, Humidity.Normal, Wind.Weak])) # Sunny + High = NO self.assertEqual( Label.NO, tree.predict( [Outlook.Sunny, Temperature.Cool, Humidity.High, Wind.Strong])) self.assertEqual( Label.NO, tree.predict( [Outlook.Sunny, Temperature.Hot, Humidity.High, Wind.Weak])) # Rain + Weak = Yes self.assertEqual( Label.YES, tree.predict( [Outlook.Rain, Temperature.Cool, Humidity.Normal, Wind.Weak])) self.assertEqual( Label.YES, tree.predict( [Outlook.Rain, Temperature.Hot, Humidity.High, Wind.Weak])) # Rain + Strong = No self.assertEqual( Label.NO, tree.predict( [Outlook.Rain, Temperature.Cool, Humidity.Normal, Wind.Strong])) self.assertEqual( Label.NO, tree.predict( [Outlook.Rain, Temperature.Hot, Humidity.High, Wind.Strong])) tree.print()