Fixed bug in ID3 and linked to the document (#32)

bfortuner · Aug 18, 2018 · 373bb9e · 373bb9e
1 parent a7679bd
commit 373bb9e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/code/id3_decision_tree_simple.py b/code/id3_decision_tree_simple.py
@@ -115,7 +115,11 @@ def _query(tree, tmp_test_array):
             assert len(tmp_test_array.shape) == 2, \
                 "Make sure your test data is 2d array"
 
-            start_node = tree[0, :]
+            if isinstance(tree,list):
+                start_node = tree # only the 1 row in data
+            else:
+                start_node = tree[0,:] # Iteratively hit first row
+
             test_feat, test_val, left_tree_jump, right_tree_jump = \
                 start_node[0], start_node[1], start_node[2], start_node[3]
 

diff --git a/docs/classification_algos.rst b/docs/classification_algos.rst
@@ -20,7 +20,7 @@ Be the first to `contribute! <https://github.com/bfortuner/ml-cheatsheet>`__
 Decision Trees
 ==============
 
-Be the first to `contribute! <https://github.com/bfortuner/ml-cheatsheet>`__
+ID3 decision tree: `code example <https://github.com/bfortuner/ml-cheatsheet/blob/master/code/id3_decision_tree_simple.py>`__
 
 K-Nearest Neighbor
 ==================