aimacode · norvig · Mar 25, 2017 · Mar 24, 2017 · Mar 24, 2017 · Mar 24, 2017
diff --git a/learning.py b/learning.py
@@ -42,9 +42,9 @@ def hamming_distance(predictions, targets):
 
 
 class DataSet:
-    """A data set for a machine learning problem.  It has the following fields:
+    """A data set for a machine learning problem. It has the following fields:
 
-    d.examples   A list of examples.  Each one is a list of attribute values.
+    d.examples   A list of examples. Each one is a list of attribute values.
     d.attrs      A list of integers to index into an example, so example[attr]
                  gives a value. Normally the same as range(len(d.examples[0])).
     d.attrnames  Optional list of mnemonic names for corresponding attrs.
@@ -60,14 +60,16 @@ class DataSet:
                  since that can handle any field types.
     d.name       Name of the data set (for output display only).
     d.source     URL or other source where the data came from.
+    d.exclude    A list of attribute indexes to exclude from d.inputs. Elements
+                 of this list can either be integers (attrs) or attrnames.
 
     Normally, you call the constructor and you're done; then you just
     access fields like d.examples and d.target and d.inputs."""
 
     def __init__(self, examples=None, attrs=None, attrnames=None, target=-1,
                  inputs=None, values=None, distance=mean_boolean_error,
                  name='', source='', exclude=()):
-        """Accepts any of DataSet's fields.  Examples can also be a
+        """Accepts any of DataSet's fields. Examples can also be a
         string or file from which to parse examples using parse_csv.
         Optional parameter: exclude, as documented in .setproblem().
         >>> DataSet(examples='1, 2, 3')
@@ -107,14 +109,14 @@ def setproblem(self, target, inputs=None, exclude=()):
         to not use in inputs. Attributes can be -n .. n, or an attrname.
         Also computes the list of possible values, if that wasn't done yet."""
         self.target = self.attrnum(target)
-        exclude = map(self.attrnum, exclude)
+        exclude = list(map(self.attrnum, exclude))
         if inputs:
             self.inputs = removeall(self.target, inputs)
         else:
             self.inputs = [a for a in self.attrs
                            if a != self.target and a not in exclude]
         if not self.values:
-            self.values = list(map(unique, zip(*self.examples)))
+            self.update_values()
         self.check_me()
 
     def check_me(self):
@@ -149,6 +151,9 @@ def attrnum(self, attr):
         else:
             return attr
 
+    def update_values(self):
+        self.values = list(map(unique, zip(*self.examples)))
+
     def sanitize(self, example):
         """Return a copy of example, with non-input attributes replaced by None."""
         return [attr_i if i in self.inputs else None
@@ -165,6 +170,7 @@ def classes_to_numbers(self,classes=None):
     def remove_examples(self,value=""):
         """Remove examples that contain given value."""
         self.examples = [x for x in self.examples if value not in x]
+        self.update_values()
 
     def __repr__(self):
         return '<DataSet({}): {:d} examples, {:d} attributes>'.format(

diff --git a/tests/test_learning.py b/tests/test_learning.py
@@ -4,6 +4,10 @@
 from utils import DataFile
 
 
+def test_exclude():
+    iris = DataSet(name='iris', exclude=[3])
+    assert iris.inputs == [0, 1, 2]
+
 
 def test_parse_csv():
     Iris = DataFile('iris.csv').read()
@@ -38,6 +42,7 @@ def test_k_nearest_neighbors():
     kNN = NearestNeighborLearner(iris,k=3)
     assert kNN([5,3,1,0.1]) == "setosa"
 
+
 def test_decision_tree_learner():
     iris = DataSet(name="iris")
 
@@ -47,21 +52,23 @@ def test_decision_tree_learner():
 
 def test_neural_network_learner():
     iris = DataSet(name="iris")
+    iris.remove_examples("virginica")
+
     classes = ["setosa","versicolor","virginica"]
-
     iris.classes_to_numbers()
 
     nNL = NeuralNetLearner(iris)
-    # NeuralNetLearner might be wrong. Just check if prediction is in range
+    # NeuralNetLearner might be wrong. Just check if prediction is in range.
     assert nNL([5,3,1,0.1]) in range(len(classes))
 
 
 def test_perceptron():
     iris = DataSet(name="iris")
+    iris.remove_examples("virginica")
+
     classes = ["setosa","versicolor","virginica"]
-
     iris.classes_to_numbers()
 
     perceptron = PerceptronLearner(iris)
-    # PerceptronLearner might be wrong. Just check if prediction is in range
+    # PerceptronLearner might be wrong. Just check if prediction is in range.
     assert perceptron([5,3,1,0.1]) in range(len(classes))