Scala中如何用map、zip将两层循环加权平均改写为函数式风格

改写后的函数式实现 #

你可以通过zip将emb和权重先绑定，再利用向量转置的思路按维度聚合，完全不需要通过下标索引，写法更简洁也更符合函数式编程习惯：

val aggAvg = (emb: Seq[Seq[Float]], weights: Seq[Float]) => {
  val weightSum = weights.sum
  emb.zip(weights)
    .map { case (vec, w) => vec.map(_ * w) } // 每个向量先乘对应权重
    .reduce((a, b) => a.zip(b).map { case (x, y) => x + y }) // 按维度求和
    .map(_ / weightSum) // 每个维度除以权重总和
}

Scala原生Seq自带transpose方法，还可以写出更直观的简化版本：

val aggAvg = (emb: Seq[Seq[Float]], weights: Seq[Float]) => {
  val weightSum = weights.sum
  emb.transpose // 按维度转置，转后每一个子序列就是同一个维度的所有取值
    .map(dimVec => dimVec.zip(weights).map { case (v, w) => v * w }.sum / weightSum)
}

实现说明 #

• 提前计算权重总和避免重复计算，相比原写法每次循环都调用weights.sum性能更优
• 两种写法都完全规避了下标访问，逻辑和原实现完全等价，输入测试用例都会返回Seq(3.4, 4.4, 5.4)的正确结果
• 第二种用transpose的写法可读性更高，转置后每一组数据直接对应要计算加权平均的维度所有取值，直接绑定权重计算即可

内容的提问来源于stack exchange，提问作者user17101610