minor change to variance calculation

apache · Jun 29, 2016 · af1ff66 · af1ff66
1 parent e8b8914
commit af1ff66
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/tree/impurity/Variance.scala b/mllib/src/main/scala/org/apache/spark/mllib/tree/impurity/Variance.scala
@@ -106,11 +106,8 @@ object Variance extends Impurity {
     val rightSumSquares = parentSumSquares - leftSumSquares
 
     val parentImpurity = (parentSumSquares - (parentSum * parentSum) / totalCount) / totalCount
-    val leftImpurity = (leftSumSquares - (leftSum * leftSum) / leftCount) / leftCount
-    val rightImpurity = (rightSumSquares - (rightSum * rightSum) / rightCount) / leftCount
-
-    val leftWeighted = leftImpurity * leftCount / totalCount
-    val rightWeighted = rightImpurity * rightCount / totalCount
+    val leftWeighted = (leftSumSquares - (leftSum * leftSum) / leftCount) / totalCount
+    val rightWeighted = (rightSumSquares - (rightSum * rightSum) / rightCount) / totalCount
     val gain = parentImpurity - leftWeighted - rightWeighted
 
     if (gain < minInfoGain) {