I've decided to implement the ID3 Decision Tree algorithm in Python based on what I've learned from George F. Luger's textbook on AI (and other secondary readings). As far as I know, the code is working correctly:
import math from random import choice from collections import defaultdict, namedtuple class TreeNode: """Used to implement the tree structure.""" def __init__(self, type = None, value = None): #type may be "property" or "leaf". self.type = type #The edges of the tree are represented as dict keys #and the nodes as the dict values. self.branches = dict() #This variable contains the name of the property represented by the node #or the label, if it's a leaf node. self.value = value class DecisionTree: """This class implements a simple Decision Tree for classification tasks. After providing a dataset, the train method should be used to create the tree. After that, the classify method may be used to classify a new example by traversing the tree.""" def __init__(self, dataset): """The dataset must be a sequence of named tuples representing each sample of the training set, which must contain the values for each attribute and also their label. Attributes might have any name but the samples must have a 'label' attribute.""" self.dataset = dataset self.tree = None self.labels = None def train(self): """Create the decision tree based on the dataset and stores it in self.tree""" if len(self.dataset) == 0: print("The dataset was not provided.") return False #Find out the list of properties (assuming all examples have the same structure) properties = list(dataset._fields) properties.remove("label") #Find out the possible labels for each sample. self.labels = set([getattr(x, "label") for x in self.dataset]) self.tree = self._induce_tree(self.dataset, properties) def classify(self, example): """Return a string (or a list of strings) containing the label (or possible labels) of the given example.""" tree = self.tree while tree.type != "leaf": tree = tree.branches[getattr(example, tree.value)] return tree.value def _calc_entropy(self, probabilities): """Return the entropy of a given random variable given the probabilities (in the form of a list/tuple) of each of its possible outcomes""" entropy = 0 for p in probabilities: entropy -= 0 if p == 0 else p * math.log(p, 2) return entropy def _calc_expected_entropy(self, partition): """Return the expected information needed to complete the tree""" #Get the total number of elements (population) total_pop = 0 for p in partition.values(): total_pop += len(p) expectation = 0 for p in partition.values(): p_pop = len(p) p_labels = [getattr(s, "label") for s in p] #Get the probability of each label ocurrency in the given #partition: No. of elements with the label divided by #the total no. of elements of the partitions. probs =  for label in self.labels: probs.append(p_labels.count(label)/p_pop) #Expected information needed = #Sigma |Partition_i| / Total population * Entropy of Partition_i #for i from 1 to n (no. of partitions). expectation += (p_pop/total_pop) * self._calc_entropy(probs) return expectation def _choose_property(self, example_set, properties): rank = list() for p in properties: #Partition the example set in groups of samples that #share the same value for the property p. partition = defaultdict(lambda: list()) for s in example_set: partition[getattr(s, p)].append(s) rank.append((p, self._calc_expected_entropy(partition), partition)) #Sort by expected entropy in increasing order. rank.sort(key = lambda x: x) #Since rank is a list of tuples containing the property, #the expected entropy and partition for a subtree with that property #rank in the sorted rank will be the property #with the least expected entropy (or with the most significant #information gain) and rank will be the partition of #the example set for that property. return (rank, rank) def _induce_tree(self, example_set, properties): """Recursive algorithm for inducing a decision tree based in ID3 algorithm as explained by George F. Luger's textbook (2009). example_set: a list named tuples containing the samples properties: a list of strings containing the remaining properties to be chosen in the tree induction""" #If all samples belong to the same class, then #produce a leaf node for that class (label). labels = set([getattr(s, "label") for s in example_set]) if len(labels) == 1: return TreeNode(type = "leaf", value = list(labels)) #Else if there are no more properties to evaluate #but there are still examples with different labels #pick one at random and produce a leaf node with it. if not properties: return TreeNode(type = "leaf", value = choice(example_set).getatrr("label")) #Else, chose a property and make it the root of a new subtree. p, partition = self._choose_property(example_set, properties) new_root = TreeNode(type = "property", value = p) new_properties = list(properties) new_properties.remove(p) #Create a new branch in the subtree for each partition. for value, block in partition.items(): new_root.branches[value] = self._induce_tree(block, new_properties) return new_root #This is just auxiliary code... def generate_pydot_graph(graph, node, edge_name = "none"): """A recursive function to convert the generated tree into a pydot graph""" node_name = edge_name + '-' + node.value if node.type == "leaf": shape = "box" else: shape = "ellipse" graph.add_node(pydot.Node(node_name, label = node.value, shape = shape)) if node.type == "property": for edge, sibling in node.branches.items(): sibling_name = edge+'-'+sibling.value graph.add_edge(pydot.Edge(node_name, sibling_name, label = edge)) generate_pydot_graph(graph, sibling, edge) if __name__ == "__main__": Sample = namedtuple("Sample", ['label', 'credit_history', 'debit', 'collateral', 'income']) dataset = list() dataset.append(Sample('high', 'bad', 'high', 'none', '0-15k')) dataset.append(Sample('high', 'unknown', 'high', 'none', '15-35k')) dataset.append(Sample('moderate', 'unknown', 'low', 'none', '15-35k')) dataset.append(Sample('high', 'unknown', 'low', 'none', '0-15k')) dataset.append(Sample('low', 'unknown', 'low', 'none', '35k+')) dataset.append(Sample('low', 'unknown', 'low', 'adequate', '35k+')) dataset.append(Sample('high', 'bad', 'low', 'none', '0-15k')) dataset.append(Sample('moderate', 'bad', 'low', 'adequate', '35k+')) dataset.append(Sample('low', 'good', 'low', 'none', '35k+')) dataset.append(Sample('low', 'good', 'high', 'adequate', '35k+')) dataset.append(Sample('high', 'good', 'high', 'none', '0-15k')) dataset.append(Sample('moderate', 'good', 'high', 'none', '15-35k')) dataset.append(Sample('low', 'good', 'high', 'none', '35k+')) dataset.append(Sample('high', 'bad', 'high', 'none', '15-35k')) decision_tree = DecisionTree(dataset) decision_tree.train() #Testing the tree... inst_fields = list(Sample._fields) inst_fields.remove("label") NewInstance = namedtuple("NewInstance", inst_fields) tests = list() tests.append(NewInstance("unknown", "low", "none", "15-35k")) #moderate tests.append(NewInstance("good", "node", "none", "0-15k")) #high tests.append(NewInstance("unknown", "high", "adequate", "15-35k")) #high for t in tests: print(decision_tree.classify(t)) #Generate a PNG image with a graphical representation of the tree #to compare with the one in Luger's book :) import pydot graph = pydot.Dot("decision_tree", graph_type = "graph") generate_pydot_graph(graph, decision_tree.tree) graph.write_png("output.png")
I'm particularly interested in feedback on the choices I've made concerning the data structures used (sequences of namedtuples for the dataset, dicts for partitions and so on). I'd also like to know if I commented and documented the code properly and if there are "more pythonic" ways of doing things.
My main goal was to implement the algorithm in a way that it would be easy to read and understand, so performance was not a primary concern of mine. Nonetheless, I'd appreciate tips on that too. Thank you!