【C++】合計が閾値以下となる正方形の最大辺の長さを求めるアルゴリズム

問題概要

m × n の行列 mat と整数 threshold(閾値)が与えられます。このとき、要素の合計が閾値以下となる正方形の中で最大の辺の長さを求めてください。条件を満たす正方形が存在しない場合は 0 を返します。

例えば、次のような入力が与えられたとします。

閾値が 4 の場合、出力は 2 となります。上の表の左上、および下の表の左下にある 2×2 の領域(青・緑色部分)は、合計が 1 + 1 + 1 + 1 = 4 となり閾値以下です。一方、3×3 の正方形の合計は 15 となり条件を満たさないため、答えは 2 になります。

解法のアプローチ

この問題は二次元累積和(プレフィックスサム)と二分探索を組み合わせることで効率的に解くことができます。全体の流れは以下の通りです。

1. 行列を二次元累積和に変換し、任意の矩形領域の合計を高速に計算できるようにする。
2. 正方形の辺の長さの候補に対して二分探索を行う。
3. 各候補の長さについて、「合計が閾値以下の正方形が存在するか」を累積和で判定する(ok 関数)。

ok 関数(判定処理)の詳細

ok(x, mat, th) は「合計が th 以下の x × x の正方形が存在するか」を返す関数です。

• curr := 0 で初期化する。
• r を x−1 から行数−1 まで、c を x−1 から列数−1 まで走査する。
  • curr := mat[r][c](右下隅の累積和)
  • c − x ≥ 0 なら、curr から mat[r][c−x] を引く(左側の余分な領域を除去)
  • r − x ≥ 0 なら、curr から mat[r−x][c] を引く(上側の余分な領域を除去)
  • c − x ≥ 0 かつ r − x ≥ 0 なら、curr に mat[r−x][c−x] を加える(二重に引いた部分を戻す)
  • curr ≤ th なら true を返す
• 最後まで見つからなければ false を返す。

メイン関数(maxSideLength)の詳細

• r := 行数、c := 列数、low := 1、high := min(r, c)、ans := 0 と初期化する。
• まず行方向に累積和を作る:i を 1 から c−1 まで、j を 0 から r−1 まで走査し、mat[j][i] += mat[j][i−1]。
• 次に列方向に累積和を作る:i を 1 から r−1 まで、j を 0 から c−1 まで走査し、mat[i][j] += mat[i−1][j]。
• low ≤ high の間、以下を繰り返す:
  • mid := low + (high − low) / 2
  • ok(mid, mat, th) が真なら ans := mid、low := mid + 1(より大きな辺を探索)
  • 偽なら high := mid − 1(より小さな辺を探索)
• ans を返す。

C++ 実装例

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

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
class Solution {
public:
   bool ok(int x, vector < vector<int> >& mat, int th){
      lli current = 0;
      for(int r = x - 1; r < mat.size(); r++){
         for(int c = x - 1; c < mat[0].size(); c++){
            current = mat[r][c];
            if(c - x >= 0)current -= mat[r][c-x];
            if(r -x >= 0)current -= mat[r - x][c];
            if(c - x >= 0 && r - x >= 0)current += mat[r-x][c-x];
            if(current <= th)return true;
         }
      }
      return false;
   }
   int maxSideLength(vector<vector<int>>& mat, int th) {
      int r = mat.size();
      int c = mat[0].size();
      int low = 1;
      int high = min(r, c);
      int ans = 0;
      for(int i = 1; i < c; i++){
         for(int j = 0; j < r; j++){
            mat[j][i] += mat[j][i - 1];
         }
      }
      for(int i = 1; i < r; i++){
         for(int j = 0; j < c; j++){
            mat[i][j] += mat[i - 1][j];
         }
      }
      while(low <= high){
         int mid = low + ( high - low ) / 2;
         if(ok(mid, mat, th)){
            ans = mid;
            low = mid + 1;
         }
         else{
            high = mid - 1;
         }
      }
      return ans;
   }
};
main(){
   vector<vector<int>> v = {{1,1,3,2,4,3,2},{1,1,3,2,4,3,2},{1,1,3,2,4,3,2}};
   Solution ob;
   cout << (ob.maxSideLength(v, 4));
}

入力

[[1,1,3,2,4,3,2],[1,1,3,2,4,3,2],[1,1,3,2,4,3,2]]
4

出力

2

計算量の評価

累積和の構築には O(mn)、二分探索の各ステップにおける判定にも O(mn) かかり、探索回数は O(log(min(m, n))) 回です。したがって、全体の時間計算量は O(mn · log min(m, n)) となります。また、元の行列をそのまま累積和行列として書き換えているため、追加の空間計算量は O(1) で済みます。