fix bug in decision stump using wrong value for number of classes

EdwardRaff · Mar 2, 2016 · 290ef4f · 290ef4f
1 parent 5412a2d
commit 290ef4f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/JSAT/src/jsat/classifiers/trees/DecisionStump.java b/JSAT/src/jsat/classifiers/trees/DecisionStump.java
@@ -394,7 +394,7 @@ public List<List<DataPointPair<Integer>>> trainC(List<DataPointPair<Integer>> da
                 aSplit = listOfLists(catAttributes[attribute].getNumOfCategories());
                 split_scores = new ImpurityScore[aSplit.size()];
                 for(int i=0; i < split_scores.length; i++)
-                    split_scores[i] = new ImpurityScore(aSplit.size(), gainMethod);
+                    split_scores[i] = new ImpurityScore(predicting.getNumOfCategories(), gainMethod);
 
                 List<DataPointPair<Integer>> wasMissing = new ArrayList<DataPointPair<Integer>>(0);
                 double missingSum = 0.0;