Problem Description

You are given a 2D integer matrix grid (0-indexed) and an integer k.

The task is to count how many submatrices meet two conditions:

1. The submatrix must contain the top-left element of the grid (element at position [0][0])
2. The sum of all elements in the submatrix must be less than or equal to k

A submatrix is defined by selecting a rectangular region from the original matrix. Since all valid submatrices must contain the top-left element grid[0][0], they all start from position [0][0] and extend to some position [i][j] where i and j represent the bottom-right corner of the submatrix.

For example, if you have a 3x3 grid, possible submatrices that contain the top-left element would include:

• Just the single element at [0][0]
• A 1x2 rectangle from [0][0] to [0][1]
• A 2x1 rectangle from [0][0] to [1][0]
• A 2x2 rectangle from [0][0] to [1][1]
• And so on, up to the entire grid

The solution uses a two-dimensional prefix sum approach where s[i][j] represents the sum of all elements in the submatrix from [0][0] to [i-1][j-1]. The prefix sum is calculated using the formula:

s[i][j] = s[i-1][j] + s[i][j-1] - s[i-1][j-1] + x

where x is the current element at position [i-1][j-1] in the original grid. Since each prefix sum represents a valid submatrix starting from the top-left corner, we simply count how many of these sums are less than or equal to k.

Intuition

The key insight is recognizing that since all valid submatrices must contain the top-left element grid[0][0], every submatrix we need to consider has a fixed starting point at [0][0] and extends to some position [i][j].

This constraint simplifies our problem significantly. Instead of considering all possible submatrices in the grid (which would require checking different starting and ending points), we only need to consider submatrices that start from the top-left corner. This means for an m x n grid, we have exactly m * n possible submatrices to check.

The natural approach would be to calculate the sum for each of these submatrices individually, but that would be inefficient. For each submatrix ending at [i][j], we'd need to sum all elements from [0][0] to [i][j], leading to repeated calculations.

This is where the prefix sum technique becomes valuable. As we iterate through the grid from top-left to bottom-right, we can build up the sum incrementally. When we're at position [i][j], we already know:

• The sum of the submatrix from [0][0] to [i-1][j] (the rectangle directly above)
• The sum of the submatrix from [0][0] to [i][j-1] (the rectangle to the left)
• The sum of the submatrix from [0][0] to [i-1][j-1] (the overlapping rectangle)

Using these previously calculated values, we can compute the sum for the submatrix ending at [i][j] by adding the current element, adding the sums from above and left, and subtracting the overlap that was counted twice: s[i][j] = s[i-1][j] + s[i][j-1] - s[i-1][j-1] + current_element.

Since each prefix sum s[i][j] represents exactly one valid submatrix (from [0][0] to [i][j]), we can simply check if this sum is less than or equal to k and increment our counter accordingly. This gives us an efficient solution that processes each element only once.

Learn more about Prefix Sum patterns.

Solution Approach

The implementation uses a two-dimensional prefix sum approach to efficiently calculate the sum of all submatrices starting from the top-left corner.

Data Structure Setup: We create a 2D array s with dimensions (m+1) x (n+1) where m and n are the dimensions of the original grid. The extra row and column (initialized to 0) serve as padding to handle boundary cases elegantly, eliminating the need for special checks when accessing s[i-1][j] or s[i][j-1].

Algorithm Implementation:

1. Initialize the prefix sum array s with all zeros and a counter ans to track valid submatrices.

2. Iterate through each element of the grid using nested loops. We use enumerate(grid, 1) and enumerate(row, 1) to get 1-indexed positions, which align with our padded prefix sum array.

3. For each position (i, j), calculate the prefix sum using the formula:

   s[i][j] = s[i-1][j] + s[i][j-1] - s[i-1][j-1] + x

   Where:

   • s[i-1][j] is the sum of the submatrix from [0][0] to [i-2][j-1]
   • s[i][j-1] is the sum of the submatrix from [0][0] to [i-1][j-2]
   • s[i-1][j-1] is the overlapping region that gets counted twice
   • x is the current element grid[i-1][j-1]

4. After calculating each prefix sum, immediately check if s[i][j] <= k. If true, increment the answer counter. This works because each s[i][j] represents exactly one submatrix starting from [0][0] and ending at [i-1][j-1] in the original grid.

Time and Space Complexity:

• Time: O(m * n) where we visit each element exactly once
• Space: O(m * n) for the prefix sum array

The elegance of this solution lies in combining the prefix sum calculation with the counting logic in a single pass, making it both efficient and concise.

Example Walkthrough

Let's walk through a concrete example to illustrate the solution approach.

Consider a 3x3 grid with k = 7:

grid = [[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]]
k = 7

We'll create a prefix sum array s with dimensions 4x4 (padded with zeros):

s = [[0, 0, 0, 0],
     [0, 0, 0, 0],
     [0, 0, 0, 0],
     [0, 0, 0, 0]]

Step-by-step calculation:

Position [0][0] (grid value = 1):

• Calculate: s[1][1] = s[0][1] + s[1][0] - s[0][0] + 1 = 0 + 0 - 0 + 1 = 1
• This represents the submatrix containing only element [0][0]
• Since 1 ≤ 7, increment counter: ans = 1

Position [0][1] (grid value = 1):

• Calculate: s[1][2] = s[0][2] + s[1][1] - s[0][1] + 1 = 0 + 1 - 0 + 1 = 2
• This represents the 1x2 submatrix from [0][0] to [0][1]
• Since 2 ≤ 7, increment counter: ans = 2

Position [0][2] (grid value = 1):

• Calculate: s[1][3] = s[0][3] + s[1][2] - s[0][2] + 1 = 0 + 2 - 0 + 1 = 3
• This represents the 1x3 submatrix from [0][0] to [0][2]
• Since 3 ≤ 7, increment counter: ans = 3

Position [1][0] (grid value = 1):

• Calculate: s[2][1] = s[1][1] + s[2][0] - s[1][0] + 1 = 1 + 0 - 0 + 1 = 2
• This represents the 2x1 submatrix from [0][0] to [1][0]
• Since 2 ≤ 7, increment counter: ans = 4

Position [1][1] (grid value = 1):

• Calculate: s[2][2] = s[1][2] + s[2][1] - s[1][1] + 1 = 2 + 2 - 1 + 1 = 4
• This represents the 2x2 submatrix from [0][0] to [1][1]
• Since 4 ≤ 7, increment counter: ans = 5

Position [1][2] (grid value = 1):

• Calculate: s[2][3] = s[1][3] + s[2][2] - s[1][2] + 1 = 3 + 4 - 2 + 1 = 6
• This represents the 2x3 submatrix from [0][0] to [1][2]
• Since 6 ≤ 7, increment counter: ans = 6

Position [2][0] (grid value = 1):

• Calculate: s[3][1] = s[2][1] + s[3][0] - s[2][0] + 1 = 2 + 0 - 0 + 1 = 3
• This represents the 3x1 submatrix from [0][0] to [2][0]
• Since 3 ≤ 7, increment counter: ans = 7

Position [2][1] (grid value = 1):

• Calculate: s[3][2] = s[2][2] + s[3][1] - s[2][1] + 1 = 4 + 3 - 2 + 1 = 6
• This represents the 3x2 submatrix from [0][0] to [2][1]
• Since 6 ≤ 7, increment counter: ans = 8

Position [2][2] (grid value = 1):

• Calculate: s[3][3] = s[2][3] + s[3][2] - s[2][2] + 1 = 6 + 6 - 4 + 1 = 9
• This represents the entire 3x3 grid
• Since 9 > 7, do not increment counter: ans = 8

Final answer: 8 submatrices have sum ≤ 7

Solution Implementation

Time and Space Complexity

The time complexity is O(m × n), where m is the number of rows and n is the number of columns of the matrix. This is because the algorithm uses two nested loops that iterate through each element of the grid exactly once. The outer loop runs m times (for each row), and the inner loop runs n times (for each column). All operations inside the loops (array access, arithmetic operations, and comparison) take constant time O(1).

The space complexity is O(m × n), where m and n are the dimensions of the input matrix. This is due to the 2D prefix sum array s that is created with dimensions (m + 1) × (n + 1). The extra row and column are added to handle boundary cases, but the space requirement is still proportional to the size of the input grid. The only additional space used is for the loop variables and the answer counter, which require O(1) space.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Index Mapping Between Grid and Prefix Sum Array

One of the most common mistakes is confusion about the index mapping between the original grid and the padded prefix sum array. Since the prefix sum array has an extra row and column for padding, prefix_sum[i][j] corresponds to the sum from grid[0][0] to grid[i-1][j-1], not grid[i][j].

Incorrect Implementation:

# Wrong: directly using grid indices for prefix_sum
prefix_sum[i][j] = (prefix_sum[i-1][j] + 
                   prefix_sum[i][j-1] - 
                   prefix_sum[i-1][j-1] + 
                   grid[i][j])  # Error: index out of bounds when i=rows or j=cols

Correct Implementation:

# Correct: adjusting indices properly
prefix_sum[i][j] = (prefix_sum[i-1][j] + 
                   prefix_sum[i][j-1] - 
                   prefix_sum[i-1][j-1] + 
                   grid[i-1][j-1])  # Correct: maps to actual grid element

2. Forgetting the Inclusion-Exclusion Principle

Another frequent error is forgetting to subtract the overlapping region prefix_sum[i-1][j-1] when calculating the prefix sum. This leads to double-counting the rectangular region that appears in both prefix_sum[i-1][j] and prefix_sum[i][j-1].

Incorrect Implementation:

# Wrong: missing subtraction of overlapping area
prefix_sum[i][j] = (prefix_sum[i-1][j] + 
                   prefix_sum[i][j-1] + 
                   grid[i-1][j-1])  # Error: double-counts the overlap

Correct Implementation:

# Correct: properly applying inclusion-exclusion
prefix_sum[i][j] = (prefix_sum[i-1][j] + 
                   prefix_sum[i][j-1] - 
                   prefix_sum[i-1][j-1] +  # Subtract the overlap
                   grid[i-1][j-1])

3. Misunderstanding the Problem Requirements

Some might incorrectly interpret the problem as finding ALL possible submatrices in the grid, not just those starting from the top-left corner [0][0].

Incorrect Approach:

# Wrong: counting all possible submatrices
count = 0
for i1 in range(rows):
    for j1 in range(cols):
        for i2 in range(i1, rows):
            for j2 in range(j1, cols):
                # Calculate sum of submatrix from (i1,j1) to (i2,j2)
                # This counts ALL submatrices, not just those from (0,0)

Correct Understanding: Since all valid submatrices must contain grid[0][0], they all start from position [0][0]. Each prefix sum prefix_sum[i][j] represents exactly one such submatrix, making the solution much simpler.

4. Off-by-One Errors in Loop Ranges

Using incorrect loop ranges can lead to missing elements or accessing out-of-bounds indices.

Incorrect Implementation:

# Wrong: loop ranges don't cover all elements
for i in range(rows):  # Should be range(1, rows + 1) for prefix_sum
    for j in range(cols):  # Should be range(1, cols + 1) for prefix_sum
        prefix_sum[i][j] = ...  # This misses the padding logic

Correct Implementation:

# Correct: proper ranges for 1-indexed prefix sum array
for i in range(1, rows + 1):
    for j in range(1, cols + 1):
        prefix_sum[i][j] = ...

These pitfalls highlight the importance of understanding the relationship between array indices, the mathematical formula for 2D prefix sums, and the specific problem constraints.