具有行主矩阵特征 C++ 的平方欧几里得距离

Question

由于我计划使用 pybind11 将 numpy 数组传递到我的 C++ 代码中，自然我想使用 Row Major 矩阵进行计算。我发现了堆栈上平方欧几里德距离的（一个线性）实现

typedef Eigen::MatrixXd Matrix;
void squared_dist(const Matrix& X1, const Matrix& X2, Matrix& D) {
    D = ((-2 * X1.transpose() * X2).colwise() + X1.colwise().squaredNorm().transpose()).rowwise() + X2.colwise().squaredNorm();
}

但这要求 X1、X2 和 D 是默认的列主矩阵。我将如何为行主要矩阵实现类似的单行？

Answer 1

您可以使用该单行函数的模板版本，以便它可以接受 RowMajor 以及 ColMajor Eigen::Matrix 参数：

template<class T>
void squared_dist(const T& X1, const T& X2, T& D) {
    D = ((-2 * X1.transpose() * X2).colwise() + X1.colwise().squaredNorm().transpose()).rowwise() + X2.colwise().squaredNorm();
}

This godbolt demo 展示了如何在代码中使用该函数。