diff --git a/decision_tree.ipynb b/decision_tree.ipynb
index 07da773..845c565 100644
--- a/decision_tree.ipynb
+++ b/decision_tree.ipynb
@@ -657,7 +657,7 @@
     "\n",
     "    for col in range(n_features):  # for each feature\n",
     "\n",
-    "        values = set([row[col] for row in rows])  # unique values in the column\n",
+    "        values = unique_vals(rows, col)  # unique values in the column\n",
     "\n",
     "        for val in values:  # for each value\n",
     "\n",
diff --git a/decision_tree.py b/decision_tree.py
index 495ac0a..572c0e8 100644
--- a/decision_tree.py
+++ b/decision_tree.py
@@ -231,7 +231,7 @@ def find_best_split(rows):
 
     for col in range(n_features):  # for each feature
 
-        values = set([row[col] for row in rows])  # unique values in the column
+        values = unique_vals(rows, col)  # unique values in the column
 
         for val in values:  # for each value