C++でマトリックスブロックの合計を求めるアルゴリズムを解説

問題概要

m × n の行列 mat と整数 K が与えられたとき、各要素 answer[i][j] が「i − K ≤ r ≤ i + K」かつ「j − K ≤ c ≤ j + K」を満たすすべての要素 mat[r][c] の合計となるような、新しい行列 answer を求めます。ただし、(r, c) は行列内の有効な位置である必要があります。

入力例

たとえば、次のような 3 × 3 の行列が与えられたとします。

ここで k = 1 の場合、出力は次のようになります。

例えば左上の要素 12 は、自身の値 1 と、範囲内に存在する右隣の 2・下隣の 4・右下の 5 を足した合計です。端のセルでは、行列の外側にはみ出す位置は計算対象から除外されます。

解法のアプローチ

この問題は、各セルごとに周囲 (2K+1) × (2K+1) の領域に含まれる要素を素直に加算していくことで解くことができます。具体的な手順は以下の通りです。

• n を行数、m を列数とします。
• n × m のサイズを持つ行列 ans を定義します。
• i を 0 から n − 1 まで繰り返します。
  • j を 0 から m − 1 まで繰り返します。
    • r を i − k から i + k まで繰り返します。
      • c を j − k から j + k まで繰り返します。
        • r と c が行列のインデックス範囲内であれば、ans[i][j] += mat[r][c] とします。
• 最後に ans を返します。

この方法の計算量は O(n × m × k²) です。行列のサイズや K の値が小さければ十分実用的ですが、入力が大きくなる場合は二次元累積和（prefix sum）を使うことで O(n × m) まで高速化できる点も覚えておくと良いでしょう。

C++による実装例

それでは、実際のコードを見て理解を深めましょう。

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<vector<auto> > v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
       cout << "[";
       for(int j = 0; j <v[i].size(); j++){
          cout << v[i][j] << ", ";
       }
       cout << "],";
   }
    cout << "]"<<endl;
}
class Solution {
public:
    vector<vector<int>> matrixBlockSum(vector<vector<int>>& mat, int k) {
       int n = mat.size();
       int m = mat[0].size();
       vector < vector <int> > ans(n , vector <int> (m));
       for(int i = 0; i < n; i++){
          for(int j = 0; j < m; j++){
             for(int r = i - k;r <= i + k; r++){
                for(int c = j - k; c <= j + k; c++){
                   if(r>= 0 && r < n && c >= 0 && c < m){
                      ans[i][j] += mat[r][c];
                   }
                }
             }
          }
       }
       return ans;
    }
};
main(){
    vector<vector<int>> v1 = {{1,2,3},{4,5,6},{7,8,9}};
    Solution ob;
    print_vector(ob.matrixBlockSum(v1, 1));
}

入力

[[1,2,3],[4,5,6],[7,8,9]]
1

出力

[[12, 21, 16],[27, 45, 33],[24, 39, 28]]

まとめ

本記事では、C++を用いてマトリックスブロックの合計を求める方法を解説しました。各セルを中心とした (2K+1) × (2K+1) の領域内の要素を単純に加算する全探索的なアプローチは直感的で実装も容易ですが、計算量は O(n × m × k²) となります。より大きな入力に対応したい場合は、二次元累積和を活用した最適化も検討してみてください。