C++で行列に配置できる1の最大数を求める方法

w × h のサイズを持つ行列 M を考えます。すべてのセルの値は 0 または 1 であり、サイズ l × l の任意の正方形部分行列に含まれる 1 の数は maxOnes 個以下でなければなりません。このとき、行列 M に含めることができる 1 の最大数を求める必要があります。

例えば、入力が w = 3、h = 3、l = 2、maxOnes = 1 の場合、出力は 4 になります。3 × 3 の行列では、どの 2 × 2 の部分行列にも 1 は 1 個までしか含められないためです。1 を 4 個配置できる最適な解は以下の通りです。

解法のアプローチ

この問題の鍵となるのは、行列 M を n × n のパターンが敷き詰められた構造として捉えることです。制約は n × n の部分行列に対して適用されるため、パターン内の各位置に 1 を置くかどうかを決めれば、行列全体の構成が自動的に決まります。なお、以下のコードでは部分行列のサイズ l を n として扱っています。

そこで、n × n のパターン内の各位置が w × h の行列全体で何回出現するかを数え、出現回数の多い位置から順に maxOnes 個を選んで 1 を配置すれば、1 の総数を最大化できます。

具体的な手順は以下の通りです。

• ret := 0 で初期化します。

• n × n のサイズの 2 次元配列 sq を作成します。

• i := 0 から開始し、i < height の間、i を 1 ずつ増やしながら以下を繰り返します。

  • j := 0 から開始し、j < width の間、j を 1 ずつ増やしながら、sq[i mod n][j mod n] の値を 1 ずつ増加させます。これにより、パターン内の各位置が行列全体で何回現れるかが記録されます。

• 配列 v を定義します。

• i := 0 から開始し、i < n の間、i を 1 ずつ増やしながら以下を繰り返します。

  • j := 0 から開始し、j < n の間、j を 1 ずつ増やしながら、sq[i][j] の値を配列 v の末尾に追加します。

• 配列 v を降順にソートします。

• i := 0、j := 0 から開始し、i < v のサイズ かつ j < maxOnes の間、i と j を 1 ずつ増やしながら、ret に v[i] を加算します。これにより、出現回数の多い位置から順に maxOnes 個を選んだ合計が得られます。

• ret を返します。

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

実装例

#include <bits/stdc++.h>
using namespace std;
class Solution {
    public:
    int maximumNumberOfOnes(int width, int height, int n, int maxOnes) {
        int ret = 0;
        vector < vector <int> > sq(n, vector <int>(n));
        for(int i = 0; i < height; i++){
            for(int j = 0; j < width; j++){
                sq[i % n][j % n]++;
            }
        }
        vector <int> v;
        for(int i = 0; i < n; i++){
            for(int j = 0; j < n ; j++){
                v.push_back(sq[i][j]);
            }
        }
        sort(v.rbegin(), v.rend());
        for(int i = 0, j = 0; i < v.size() && j < maxOnes; i++, j++){
            ret += v[i];
        }
        return ret;
    }
};
main(){
    Solution ob;
    cout << (ob.maximumNumberOfOnes(3,3,2,1));
}

入力

3, 3, 2, 1

出力

4