如何在Eigen中以无循环方式缩放数组正元素至目标总和？

我有一个Eigen::Array<float, N, 1>类型的数组，希望将其中所有正元素乘以一个正标量s，让这些正元素的总和变为目标值t。目前我已经用带循环的代码实现了这个功能：

template<int N>
void positive_scale(Eigen::Array<float, N, 1>& A, float t) {
  float pos_sum = 0;
  for (int i = 0; i < N; ++i) {
    if (A(i) > 0) pos_sum += A(i);
  }

  float s = t / pos_sum;
  for (int i = 0; i < N; ++i) {
    if (A(i) > 0) A(i) *= s;
  }
}

我想知道有没有更优雅的无循环实现方式，就像Python NumPy里这种简洁的写法：

def positive_scale(A, t):
  A[A > 0] *= t / sum(A[A > 0])

无循环实现方案 #

Eigen支持类似NumPy的向量化操作，无需手动编写循环，直接利用布尔掩码和数组块操作就能实现：

template<int N>
void positive_scale(Eigen::Array<float, N, 1>& A, float t) {
    auto mask = A > 0;
    float pos_sum = A(mask).sum();
    A(mask) *= t / pos_sum;
}

代码说明： #

• A > 0：生成与原数组同尺寸的布尔数组，标记每个位置的元素是否为正；
• A(mask)：返回原数组中满足掩码条件的元素视图（无额外内存拷贝，属于原地操作的引用）；
• .sum()：直接计算正元素的总和；
• A(mask) *= ...：对所有正元素原地应用缩放因子，完成总和调整。

这种写法和NumPy风格一致，既简洁又能充分利用Eigen的向量化优化，效率不低于手动循环，在大数组场景下表现更优。

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