Pythonで合計がk以下となる最大の長方形の合計を求めるプログラム

問題の概要

2次元行列とある値 k が与えられたとき、「要素の合計が k 以下」となる長方形領域の中で、合計が最大になるものを見つける問題を考えます。

例として、次のような行列を入力とした場合をみてみましょう。

k = 15 のとき、出力は 12 になります。これは、縦方向に [5, 7] の長方形を選ぶことで合計が 12 となり、15 以下の条件を満たす中で最大の値を得られるためです。

解法のアルゴリズム

この問題は、行の組み合わせを固定しながら累積和と集合を活用することで効率的に解けます。手順は以下の通りです。

• n を行列 a の行数、m を列数とします
• ans を十分に小さい値(事実上の負の無限大)で初期化します
• i1 を 0 から n-1 まで繰り返します
  • row をサイズ m のゼロで初期化されたリストとして用意します
  • i2 を i1 から n-1 まで繰り返します
    • j について row[j] に a[i2][j] を加算し、i1 行目から i2 行目までの各列の累積和を作ります
    • s を新しい集合とし、0 を挿入します
    • sum を 0 に初期化します
    • j を 0 から m-1 まで繰り返します
      • sum に row[j] を加算します
      • temp を、集合 s の中で (sum − k) より大きい要素すべてのリストとします
      • temp のサイズが 0 より大きい場合
        • u を temp の最小値とします
        • ans を ans と (sum − u) のうち大きい方に更新します
      • sum を s に挿入します
• 最後に ans を返します

この手法のポイントは、部分矩形の合計を「列方向の累積和 + 行方向の累積和」に分解し、(現在の累積和 − 過去の累積和) ≤ k となる過去の累積和のうち最大のものを探す点にあります。これにより、合計が k を超えない範囲で可能な限り大きな長方形の合計を効率よく求められます。なお、実運用では集合の代わりに二分探索可能なソート済み構造体(bisect や SortedList など)を使うと、探索を高速化できます。

実装例

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

class Solution:
    def solve(self, a, k):
        n = len(a)
        if n == 0:
            return 0;
        m = len(a[0])
        ans = -999999;
        for i1 in range(n):
            row = [0]*m;
            for i2 in range(i1, n):
                for j in range(m):
                    row[j] += a[i2][j]
                s = set()
                s.add(0)
                sum = 0
                for j in range(m):
                    sum += row[j];
                    temp = [e for e in s if e > (sum - k)]
                if len(temp) > 0:
                    u = min(temp)
                    ans = max(ans, sum - u)
                s.add(sum)
        return ans
ob = Solution()
matrix = [
    [5, -2],
    [7, 10]
]
k = 15
print(ob.solve(matrix, k))

入力

[
[5, -2],
[7, 10]
], 15

出力

12