【C++】スパース行列の乗算を効率的に求めるアルゴリズムと実装

2つの行列 A と B が与えられたとき、その積 AB を求めることを考えます。ここで、A の列数と B の行数は等しいものと仮定します。

例として、入力が [[1,0,0],[-1,0,3]] と [[7,0,0],[0,0,0],[0,0,1]] の場合を見てみましょう。

この場合、出力は [[7,0,0],[-7,0,3]] となります。

スパース行列の乗算を効率化するポイント

通常の行列積では O(r1 × c1 × c2) の計算量が必要ですが、スパース行列(大部分の要素が 0 である行列)の場合、0 を含む掛け算は結果に寄与しないため、そのような計算はスキップできます。

そこで本手法では、行列 A の各行について「非ゼロ要素の列インデックスと値」のペアを事前に記録しておきます。こうすることで、非ゼロ要素どうしの演算だけを実行でき、計算量を大幅に削減できます。

解法の手順

• r1 := A の行数、r2 := B の行数とします

• c1 := A の列数、c2 := B の列数とします

• 結果を格納するための r1 × c2 の2次元配列 ret を定義します

• (列インデックス, 値) のペアを格納する配列 sparseA[r1] を定義します

• i = 0 から r1 未満の間、以下を繰り返します

  • j = 0 から c1 未満の間、以下を繰り返します

    • A[i][j] が 0 でない場合、sparseA[i] の末尾に { j, A[i][j] } を追加します

• i = 0 から r1 未満の間、以下を繰り返します

  • j = 0 から sparseA[i] のサイズ未満の間、以下を繰り返します

    • k = 0 から c2 未満の間、以下を繰り返します

      • x := sparseA[i][j] の first(列インデックス)

      • B[x][k] が 0 でない場合、ret[i][k] += sparseA[i][j] の second(値)× B[x][k] とします

• ret を返します

C++による実装例

理解を深めるために、以下の実装を見てみましょう。

class Solution {
public:
    vector<vector<int>> multiply(vector<vector<int>>& A, vector<vector<int>>& B) {
        int r1 = A.size();
        int r2 = B.size();
        int c1 = A[0].size();
        int c2 = B[0].size();
        vector<vector<int>> ret(r1, vector<int>(c2));
        vector<vector<pair<int, int>> sparseA(r1);
        for(int i = 0; i < r1; i++){
            for(int j = 0; j < c1; j++){
                if(A[i][j] != 0)sparseA[i].push_back({j, A[i][j]});
            }
        }
        for(int i = 0; i < r1; i++){
            for(int j = 0; j < sparseA[i].size(); j++){
                for(int k = 0; k < c2; k++){
                    int x = sparseA[i][j].first;
                    if(B[x][k] != 0){
                        ret[i][k] += sparseA[i][j].second * B[x][k];
                    }
                }
            }
        }
        return ret;
    }
};

入力

{{1,0,0},{-1,0,3}},{{7,0,0},{0,0,0},{0,0,1}}

出力

[[7, 0, 0], [-7, 0, 3]]

計算量

前処理で A の全要素を走査するのに O(r1 × c1)、乗算本体は A の非ゼロ要素数を nnz(A) とすると O(nnz(A) × c2) で完了します。行列が疎であるほど、素朴な三重ループ O(r1 × c1 × c2) と比べて高速に動作します。