update argmax impl

rekhajoshm · Jul 17, 2015 · 127dec5 · 127dec5
1 parent 2ea6a55
commit 127dec5
Showing 1 changed file with 29 additions and 7 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/linalg/Vectors.scala b/mllib/src/main/scala/org/apache/spark/mllib/linalg/Vectors.scala
@@ -724,22 +724,44 @@ class SparseVector(
     if (size == 0) {
       -1
     } else {
+      // Find the max active entry.
       var maxIdx = indices(0)
       var maxValue = values(0)
-
-      foreachActive { (i, v) =>
+      var maxJ = 0
+      var j = 1
+      val na = numActives
+      while (j < na) {
+        val v = values(j)
         if (v > maxValue) {
-          maxIdx = i
           maxValue = v
+          maxIdx = indices(j)
+          maxJ = j
         }
+        j += 1
       }
 
-      var k = 0
-      while (k < indices.length && indices(k) == k && values(k) != 0.0) {
-        k += 1
+      // If the max active entry is nonpositive and there exists inactive ones, find the first zero.
+      if (maxValue <= 0.0 && na < size) {
+        if (maxValue == 0.0) {
+          // If there exists an inactive entry before maxIdx, find it and return its index.
+          if (maxJ < maxIdx) {
+            var k = 0
+            while (k < maxJ && indices(k) == k) {
+              k += 1
+            }
+            maxIdx = k
+          }
+        } else {
+          // If the max active value is negative, find and return the first inactive index.
+          var k = 0
+          while (k < na && indices(k) == k) {
+            k += 1
+          }
+          maxIdx = k
+        }
       }
 
-      if (maxValue <= 0.0 || k >= maxIdx) k else maxIdx
+      maxIdx
     }
   }
 }