Decision trees modified

samridh90 · Oct 21, 2012 · eb24aa9 · eb24aa9
1 parent 3a18cfc
commit eb24aa9
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/dTree/tree.py b/dTree/tree.py
@@ -79,3 +79,16 @@ def createTree(dataSet, labels):
         subLabels = labels[:]
         myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value), subLabels)
     return myTree
+
+
+def classify(inputTree, featlabels, testVec):
+    firstStr = inputTree.keys()[0]
+    secondDict = inputTree[firstStr]
+    featIndex = featlabels.index(firstStr)
+    for key in secondDict.keys():
+        if testVec[featIndex] == key:
+            if type(secondDict[key]).__name__ == "dict":
+                classLabel = classify(secondDict[key], featlabels, testVec)
+            else:
+                classLabel = secondDict[key]
+    return classLabel