Scala:从元组数组/ RDD中获取第n个元素的总和

时间:2016-01-29 18:17:54

标签: scala collections functional-programming

我有一个tuple数组,如下所示:

val a = Array((1,2,3), (2,3,4))

我想为下面的方法编写泛型方法:

def sum2nd(aa: Array[(Int, Int, Int)]) = {
      aa.map { a => a._2 }.sum
      }

所以我正在寻找一种方法:

def sumNth(aa: Array[(Int, Int, Int)], n: Int)

3 个答案:

答案 0 :(得分:9)

有几种方法可以解决这个问题。最简单的方法是使用productElement

def unsafeSumNth[P <: Product](xs: Seq[P], n: Int): Int =
  xs.map(_.productElement(n).asInstanceOf[Int]).sum

然后(注意索引从零开始,所以n = 1给我们第二个元素):

scala> val a = Array((1, 2, 3), (2, 3, 4))
a: Array[(Int, Int, Int)] = Array((1,2,3), (2,3,4))

scala> unsafeSumNth(a, 1)
res0: Int = 5

但是,这种实现可能会在运行时以两种不同的方式崩溃:

scala> unsafeSumNth(List((1, 2), (2, 3)), 3)
java.lang.IndexOutOfBoundsException: 3
  at ...

scala> unsafeSumNth(List((1, "a"), (2, "b")), 1)
java.lang.ClassCastException: java.lang.String cannot be cast to java.lang.Integer
  at ...

即,如果元组没有足够的元素,或者您要求的元素不是Int

您可以编写一个在运行时不会崩溃的版本:

import scala.util.Try

def saferSumNth[P <: Product](xs: Seq[P], n: Int): Try[Int] = Try(
  xs.map(_.productElement(n).asInstanceOf[Int]).sum
)

然后:

scala> saferSumNth(a, 1)
res4: scala.util.Try[Int] = Success(5)

scala> saferSumNth(List((1, 2), (2, 3)), 3)
res5: scala.util.Try[Int] = Failure(java.lang.IndexOutOfBoundsException: 3)

scala> saferSumNth(List((1, "a"), (2, "b")), 1)
res6: scala.util.Try[Int] = Failure(java.lang.ClassCastException: ...

这是一种改进,因为它迫使呼叫者解决失败的可能性,但它也有点烦人,因为它迫使呼叫者解决失败的可能性。

如果您愿意使用Shapeless,您可以充分利用这两个方面:

import shapeless._, shapeless.ops.tuple.At

def sumNth[P <: Product](xs: Seq[P], n: Nat)(implicit
  atN: At.Aux[P, n.N, Int]
): Int = xs.map(p => atN(p)).sum

然后:

scala> sumNth(a, 1)
res7: Int = 5

但是糟糕的人甚至不会编译:

scala> sumNth(List((1, 2), (2, 3)), 3)
<console>:17: error: could not find implicit value for parameter atN: ...

但这仍然不完美,因为这意味着第二个参数必须是一个字面数字(因为它需要在编译时知道):

scala> val x = 1
x: Int = 1

scala> sumNth(a, x)
<console>:19: error: Expression x does not evaluate to a non-negative Int literal
       sumNth(a, x)
                 ^

在许多情况下,这不是问题。

总结一下:如果您愿意承担合理代码崩溃程序的责任,请使用productElement。如果您想要更安全一些(以一些不便为代价),请将productElementTry一起使用。如果您想要编译时安全(但有一些限制),请使用Shapeless。

答案 1 :(得分:2)

你可以做这样的事情,虽然它不是真正的类型安全:

  def sumNth(aa: Array[Product], n: Int)= {
    aa.map { a =>
      a.productElement(n) match {
        case i:Int => i
        case _ => 0
      }
    }.sum
  }

sumNth(Array((1,2,3), (2,3,4)), 2) // 7

答案 2 :(得分:2)

不使用shapeless的Antoher类型安全方法是提供一个函数来提取所需的元素:

def sumNth[T, E: Numeric](array: Array[T])(extract: T => E) =
  array.map(extract).sum

然后你可以像这样定义sum2nd

def sum2nd(array: Array[(Int, Int, Int)]): Int = sumNth(array)(_._2)

或者像这样:

val sum2nd: Array[(Int, Int, Int)] => Int = sumNth(_)(_._2)