initial work on adding argmax to Vector and SparseVector

Author: George

mllib/src/main/scala/org/apache/spark/mllib/linalg/Vectors.scala

@@ -150,6 +150,12 @@ sealed trait Vector extends Serializable {
       toDense
     }
   }
+
+  /**
+   * Find the index of a maximal element.  Returns the first maximal element in case of a tie.
+   * Returns -1 if vector has length 0.
+   */
+  def argmax: Int
 }
 
 /**
@@ -588,11 +594,7 @@ class DenseVector(val values: Array[Double]) extends Vector {
     new SparseVector(size, ii, vv)
   }
 
-  /**
-   * Find the index of a maximal element.  Returns the first maximal element in case of a tie.
-   * Returns -1 if vector has length 0.
-   */
-  private[spark] def argmax: Int = {
+  def argmax: Int = {
     if (size == 0) {
       -1
     } else {
@@ -717,6 +719,23 @@ class SparseVector(
       new SparseVector(size, ii, vv)
     }
   }
+
+  override def argmax: Int = {
+    if (size == 0) {
+      -1
+    } else {
+      var maxIdx = 0
+      var maxValue = values(0)
+      var i = 1
+      foreachActive{ (i, v) =>
+        if(v > maxValue) {
+          maxIdx = i
+          maxValue = v
+        }
+      }
+      maxIdx
+    }
+  }
 }
 
 object SparseVector {

mllib/src/test/scala/org/apache/spark/mllib/linalg/VectorsSuite.scala

@@ -42,6 +42,7 @@ class VectorsSuite extends FunSuite {
    val vec = Vectors.dense(arr).asInstanceOf[DenseVector]
     assert(vec.size === arr.length)
     assert(vec.values.eq(arr))
+    vec.argmax
   }
 
   test("sparse vector construction") {
@@ -56,6 +57,8 @@ class VectorsSuite extends FunSuite {
     assert(vec.size === n)
     assert(vec.indices === indices)
     assert(vec.values === values)
+    val vec2 = Vectors.sparse(5,Array(0,3),values).asInstanceOf[SparseVector]
+    vec2.foreachActive( (i, v) => println(i,v))
   }
 
   test("dense to array") {