363. Max Sum of Rectangle No Larger Than K

Python

from math import inf
from bisect import insort, bisect_right


class Solution:
    def maxSumSubmatrix(self, matrix: List[List[int]], k: int) -> int:
        m, n = len(matrix), len(matrix[0])
        prefix = self.build_prefix(matrix)

        # Binary Search code refer to
        # https://leetcode.com/problems/max-sum-of-rectangle-no-larger-than-k/discuss/2490502/Clean-Python3-W-Comments-or-Prefix-Sum-and-Bisect-or-Faster-than-97
        max_sum_lk = float('-inf')
        for row1 in range(m):
            for row2 in range(row1, m):
                left_sums = [0]
                for end_col in range(n):
                    sum_here = prefix[row2][end_col] - prefix[row1 - 1][end_col]

                    pivot = bisect_right(left_sums, sum_here - k)
                    if pivot > 0 and left_sums[pivot-1] == sum_here - k:
                        return k
                    elif pivot <= end_col:
                        max_sum_lk = max(max_sum_lk, sum_here - left_sums[pivot])

                    insort(left_sums, sum_here)

        return max_sum_lk

    @staticmethod
    def build_prefix(matrix):
        m, n = len(matrix), len(matrix[0])
        prefix = [[0]*(n+1) for _ in range(m+1)]

        for i in range(m):
            for j in range(n):
                if j-1 >= 0 and i-1 >= 0:
                    prefix[i][j] = prefix[i-1][j] + prefix[i][j-1] - prefix[i-1][j-1] \
                                 + matrix[i][j]
                elif i-1 >= 0:
                    prefix[i][j] = prefix[i-1][j] + matrix[i][j]
                elif j-1 >= 0:
                    prefix[i][j] = prefix[i][j-1] + matrix[i][j]
                else:
                    prefix[i][j] = matrix[i][j]
        return prefix