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()