C++で合計がkで割り切れる部分行列を数える方法

行×列のサイズを持つ行列が入力として与えられたとき、matrix[row][col] 内に存在するすべての部分行列のうち、その要素の合計が整数 k で割り切れるものを数えるのが本記事の目標です。

例えば、行列が mat[3][3]、k が 4 の場合、該当する部分行列は以下のように表されます。

具体例で理解しよう

入力 - matrix[3][3] = { {1,1,1}, {2,2,2}, {3,3,3} }　　k=4

出力 - 合計が 'k' で割り切れる部分行列の数：4

説明 - 該当する部分行列は上図のように示されます。

入力 - matrix[3][3] = { {1,1,1}, {2,2,2}, {3,3,3} }　　k=12

出力 - 合計が 'k' で割り切れる部分行列の数：4

説明 - 該当する部分行列は下図のように示されます。

プログラムで使用するアプローチ

このアプローチでは、行列を左から右へ走査し、左端の列と右端の列の各ペアについて、部分行列の要素を配列 arr[] に追加していきます。そして、その要素の合計と、合計が k で割り切れる部分配列の個数をそれぞれ計算します。

関数 check_val() は、部分行列の要素を 1 次元配列として受け取ります。ここで累積和を計算し、k で割った余りを求め、その余りの出現回数を配列 arr_2[] に記録します。

• matrix[row][col] と整数 k を入力として受け取ります。
• 関数 check_val(int arr[], int size, int k) は、部分行列の要素からなる arr[] を受け取り、arr 内の合計が k で割り切れるすべての部分配列の個数を返します。
• 変数 count と temp を 0 で初期化します。
• 累積和を k で割った余りの出現回数を格納するための配列 arr_2[] を用意します。
• i=0 から i<size まで for ループで累積和を計算します。各 arr[i] を temp に加算し、arr_2[((temp % k) + k) % k]++ によって余りの出現回数を増やします（負の合計に対応するため、2 回 mod を取ります）。
• 再び for ループで出現回数配列 arr_2[] を走査し、値が 1 より大きい要素ごとに、count へ arr_2[i] * (arr_2[i] - 1)) / 2 を加算します。これは組み合わせ可能なすべての部分配列の個数に相当します。
• 最後に、arr_2[0] の値を count に加算します。
• 関数 matrix_divisible(int matrix[row][col], int size, int k) は入力行列を受け取り、合計が k で割り切れるすべての部分行列の個数を返します。
• 初期カウントを 0 とします。
• 一時配列 arr[size] を用意します。
• 2 重の for ループで左端の列インデックス i と右端の列インデックス j を設定します。
• 要素の合計を arr[temp] += matrix[temp][j] として計算します。
• arr[] 内の部分配列の個数分、check_val(arr, size, k) の戻り値を count に加算します。
• すべてのループが終了したら、結果として count を返します。

サンプルコード

#include <bits/stdc++.h>
using namespace std;
#define row 10
#define col 10

int check_val(int arr[], int size, int k) {
    int count = 0;
    int temp = 0;
    int arr_2[k];
    memset(arr_2, 0, sizeof(arr_2));

    for (int i = 0; i < size; i++) {
        temp = temp + arr[i];
        arr_2[((temp % k) + k) % k]++;
    }
    for (int i = 0; i < k; i++) {
        if (arr_2[i] > 1) {
            count += (arr_2[i] * (arr_2[i] - 1)) / 2;
        }
    }
    count = count + arr_2[0];
    return count;
}

int matrix_divisible(int matrix[row][col], int size, int k) {
    int count = 0;
    int arr[size];

    for (int i = 0; i < size; i++) {
        memset(arr, 0, sizeof(arr));
        for (int j = i; j < size; j++) {
            for (int temp = 0; temp < size; ++temp) {
                arr[temp] += matrix[temp][j];
            }
            count = count + check_val(arr, size, k);
        }
    }
    return count;
}
int main() {
    int matrix[row][col] = {{2,4,-1},{6,1,-9},{2,2, 1}};
    int size = 3, k = 4;
    cout << "Count of sub-matrices having sum divisible 'k' are: " << matrix_divisible(matrix, size, k);
    return 0;
}

上記のコードを実行すると、以下の出力が得られます。

出力

Count of sub-matrices having sum divisible 'k' are: 7