Added argmax to sparse vector and added unit test

feynmanliang · May 12, 2015 · df9538a · df9538a
1 parent 3cffed4
commit df9538a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 8 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
@@ -728,18 +728,11 @@ class SparseVector(
       var maxValue = values(0)
 
       foreachActive { (i, v) =>
-        if(values(i) > maxValue){
+        if(v > maxValue){
           maxIdx = i
           maxValue = v
         }
       }
-//      while(i < this.indices.size){
-//        if(values(i) > maxValue){
-//          maxIdx = indices(i)
-//          maxValue = values(i)
-//        }
-//        i += 1
-//      }
       maxIdx
     }
   }

diff --git a/mllib/src/test/scala/org/apache/spark/mllib/linalg/VectorsSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/linalg/VectorsSuite.scala
@@ -86,6 +86,10 @@ class VectorsSuite extends FunSuite {
     val vec2 = Vectors.sparse(n,indices,values).asInstanceOf[SparseVector]
     val max = vec2.argmax
     assert(max === 3)
+
+    val vec3 = Vectors.sparse(5,Array(1,3,4),Array(1.0,.5,.7))
+    val max2 = vec3.argmax
+    assert(max2 === 1)
   }
 
   test("vector equals") {