Use map and zip to be more func style in 2 for loops

Question

I implemented the following code to calculate weighted avg with for loops, how can I be more func programming style and use map and zip?

val aggAvg = (emb: Seq[Seq[Float]], weights: Seq[Float]) => {
      val embSize = emb.head.size
      val len = emb.size
      (0 until embSize)
        .map { i =>
          (0 until len).map { j =>
            emb(j)(i) * weights(j)
          }.sum / weights.sum
        }
    }

Example: Given

val emb: Seq[Seq[Float]] = Seq(Seq(1,2,3), Seq(4,5,6))
val weights: Seq[Float] = Seq(2, 8)

the output would be Seq(3.4, 4.4, 5.4) because (1 * 2 + 4 * 8) / (2 + 8) = 3.4 and so on.

Answer 1

Here is one way, although I'm not sure if it's the most elegant

val aggAvg = (emb: Seq[Seq[Float]], weights: Seq[Float]) =>
  emb.transpose.map((weights, _).zipped.map(_ * _).sum).map(_ / weights.sum)
res0: Seq[Float] = List(3.4, 4.4, 5.4)

Answer 2

This uses an Iterator in place of a zip op.

val aggAvg = (emb: Seq[Seq[Float]], weights: Seq[Float]) => {
  val wtItr = weights.iterator
  val wtSum = weights.sum
  emb.map(_.map(wtItr.next().*)).transpose.map(_.sum/wtSum)
}

Note: This can be a bit dangerous if not all of the row lengths match the weights length, but your posted aggAvg code was the same in this reguard so I left out the safety checks.