Eigen稀疏矩阵非连续块操作的高效实现方案咨询

需求说明 #

在Matlab中对30k×30k的大型稀疏矩阵执行如下操作：将a索引对应的行与b索引对应的行逐列相加，结果存回a行，最后删除b索引的所有行。Matlab代码如下：

a=[1,3,5,...];
b=[2,4,6,...];
A(a,:)=A(a,:)+A(b,:);
A(b,:)=[]

当前实现的问题 #

使用igl库的slice和slice_into函数实现时，会产生大量中间子矩阵拷贝，对于超大型稀疏矩阵来说耗时极高。现有C++代码如下：

SparseMatrix<double> A,A1, A2,A3;
igl::slice(A, a, 1, A1);
igl::slice(A, b, 1,A2);
A3 = A1 + A2;

igl::slice_into(A3, index_u, 1, A);

VectorXd A_rows_index = VectorXd::LinSpaced(A.rows(), 0, A.rows() - 1);
VectorXd A_rows_index_after(A_rows_index.size());
auto it = std::set_difference(A_rows_index.data(), A_rows_index.data() + A_rows_index.size(),
    b.data(), b.data() + b.size(),
    A_rows_index_after.data());
A_rows_index_after.conservativeResize(std::distance(A_rows_index_after.data(), it)); // resize the result
igl::slice(A, A_rows_index_after, 1, A);
A.resize();

高效实现方案 #

直接操作Eigen稀疏矩阵的底层存储，避免中间矩阵拷贝，核心思路是通过行映射关系，直接构建最终的稀疏矩阵。

核心思路 #

1. 建立原行号到新行号的映射：标记b中的行需要合并到a对应的行，同时给保留的行分配新的连续索引。
2. 遍历原矩阵的所有非零元素，根据映射关系将元素分配到新矩阵的对应位置：b行的元素直接累加到a行的新索引位置，其他行元素直接映射到新索引。
3. 利用Eigen的setFromTriplets自动合并相同(行,列)位置的元素，完成行相加操作。

代码实现 #

#include <Eigen/Sparse>
#include <vector>
#include <unordered_map>

using namespace Eigen;

void mergeAndDeleteSparseRows(SparseMatrix<double>& A, const VectorXi& a, const VectorXi& b) {
    const int origRowCount = A.rows();
    std::unordered_map<int, int> bToAMap; // 存储b行对应的a行索引
    std::vector<bool> isRowDeleted(origRowCount, false);

    // 标记待删除的行，并记录b行对应的目标a行
    for (int i = 0; i < a.size(); ++i) {
        const int bRow = b[i];
        bToAMap[bRow] = a[i];
        isRowDeleted[bRow] = true;
    }

    // 为保留的行分配新的连续索引
    std::vector<int> origToNewRow(origRowCount, -1);
    int newRowIdx = 0;
    for (int i = 0; i < origRowCount; ++i) {
        if (!isRowDeleted[i]) {
            origToNewRow[i] = newRowIdx++;
        }
    }

    // 收集所有非零元素的Triplet，准备构建新矩阵
    std::vector<Triplet<double>> triplets;
    triplets.reserve(A.nonZeros()); // 预分配内存，减少动态扩容开销

    // 遍历原矩阵的每一列
    for (int col = 0; col < A.cols(); ++col) {
        for (SparseMatrix<double>::InnerIterator it(A, col); it; ++it) {
            const int origRow = it.row();
            const double val = it.value();

            if (isRowDeleted[origRow]) {
                // 待删除行：将值合并到对应的a行
                const int targetARow = bToAMap[origRow];
                triplets.emplace_back(origToNewRow[targetARow], col, val);
            } else {
                // 保留行：直接映射到新索引
                triplets.emplace_back(origToNewRow[origRow], col, val);
            }
        }
    }

    // 构建新矩阵并与原矩阵交换（避免深拷贝）
    SparseMatrix<double> newMatrix(newRowIdx, A.cols());
    newMatrix.setFromTriplets(triplets.begin(), triplets.end());
    A.swap(newMatrix);
}

优化细节 #

• 如果a和b是有序数组，可以用普通数组代替unordered_map，进一步提升查找速度。
• 预先调用reserve为triplets分配足够内存，避免多次内存分配和拷贝。
• 若矩阵是RowMajor存储，可以转置为ColumnMajor后处理，完成后再转置回去，利用Eigen对ColumnMajor的遍历优化。

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