Fixed the sorting of support mask scores. HT: John Liu. Fixes jpmml/j…

…pmml-evaluator#82, fixes jpmml/jpmml-evaluator#93, fixes jpmml/sklearn2pmml#68, fixes jpmml/sklearn2pmml#69
jpmml · Dec 27, 2017 · d4d2af1 · d4d2af1
1 parent e26913f
commit d4d2af1
Show file tree

Hide file tree

Showing 6 changed files with 2,011 additions and 2,007 deletions.
diff --git a/src/main/java/sklearn/feature_selection/SelectKBest.java b/src/main/java/sklearn/feature_selection/SelectKBest.java
@@ -54,17 +54,22 @@ public List<Boolean> getSupportMask(){
 		for(int i = 0; i < scores.size(); i++){
 			Number score = scores.get(i);
 
-			Entry<Integer> entry = new Entry<>(i, score.doubleValue());
+			double doubleScore = score.doubleValue();
+			if(Double.isNaN(doubleScore)){
+				doubleScore = -Double.MAX_VALUE;
+			}
+
+			Entry<Integer> entry = new Entry<>(i, doubleScore);
 
 			entries.add(entry);
 		}
 
-		Collections.sort(entries, Collections.reverseOrder());
+		Collections.sort(entries);
 
 		boolean[] result = new boolean[scores.size()];
 
 		for(int i = 0, max = ValueUtil.asInt((Number)k); i < max; i++){
-			Entry<Integer> entry = entries.get(i);
+			Entry<Integer> entry = entries.get(entries.size() - (i + 1));
 
 			result[entry.getId()] = true;
 		}