LESSON 241 OF 415 · PARTIAL SUMS #2

Maximum Sum Submatrix in Swift: 2D Prefix Sums in O(m²n²)

Practice O(n² * m²) Read the full explanation

Given a m x n grid filled with integers, find the submatrix with maximum sum among all submatrices.

Example:

Input:
[
  [ 1,  -9, -10,  1],
  [-1,  10,  10,  1],
  [ 0,   9,   9, -9],
  [-1,  -1,  -1, -1]
]
Output: 38
Explanation: Submatrix [[10, 10], [9, 9]]

STEP 1 OF 17

Guided

Ctrl + Enter to check

SWIFT · PARTIAL SUMS #2

Maximum Sum Submatrix in Swift, explained

Problem overview

The method maxSum receives a grid of integers, an m by n array of arrays. A submatrix is a rectangle of neighbouring cells: a block of consecutive rows and consecutive columns. The method returns the largest sum that any submatrix of the grid can have.

In the example grid

 1  -9 -10   1
-1  10  10   1
 0   9   9  -9
-1  -1  -1  -1

the block with rows 1 to 2 and columns 1 to 2 holds 10, 10, 9, 9, which sums to 38. No other rectangle does better.

Approach: any rectangle sum from four lookups

A rectangle is fixed by two corners: the top-left (x1, y1) and the bottom-right (x2, y2). Looping over all of them is four nested loops. If the sum of each rectangle had to be computed by adding its cells, there would be two more loops inside, which is far too slow. The fix is a 2D version of prefix sums.

Let partialSums[i][j] be the sum of the rectangle from cell (0, 0) down to cell (i, j). Each entry can be built from three earlier ones by inclusion and exclusion. The sum above it and the sum to its left both contain the area above-left, so that area is counted twice and one copy is subtracted:

partialSums[i][j] = grid[i][j] + partialSums[i-1][j] + partialSums[i][j-1] - partialSums[i-1][j-1]

Reading a rectangle back works the same way, run in reverse. Start with the sum up to the bottom-right corner, remove the rows above the rectangle and the columns to its left, and add back the corner that was removed twice:

sum = P[x2][y2] - P[x1-1][y2] - P[x2][y1-1] + P[x1-1][y1-1]

Here P stands for partialSums. In the example, take x1 = 1, y1 = 1, x2 = 2, y2 = 2. The table has P[2][2] = 19, P[0][2] = -18, P[2][0] = 0 and P[0][0] = 1. The sum is 19 - (-18) - 0 + 1 = 38.

When a corner lies on row 0 or column 0, the formula needs P[-1][...] or P[...][-1]. Those cells are outside the grid, so a helper returns 0 for them, which is the right value for an empty region.

Step by step

Create result from grid[0][0]. A single cell is a submatrix, so this is a valid starting value even when every number is negative.
Create partialSums as a grid of zeros with the same size as grid, using Array(repeating:count:) for both the rows and the cells in a row. It is a var because the cells change.
Loop over the rows with i, and inside it over the columns with j.
Write the helper get above maxSum. It takes a grid and two indices, all unlabeled, and returns an Int: the cell's value, or 0 if the indices fall outside the grid.
Check both indices against 0 and against the row and column counts. When they are inside, return the cell. After the if, return 0.
Fill each prefix cell with the formula above, reading the three neighbours through get, so the first row and column need no special cases.
Choose the top-left corner with loops over x1 and y1.
Choose the bottom-right corner with loops over x2 from x1 and y2 from y1. The bottom row cannot be above the top row, and the right column cannot be left of the left one.
Compute the rectangle's sum with the four lookups and store it in a constant, sum.
Update result with the larger of its value and sum.
Return result after all the loops.

The solution in Swift

class Solution {
    func get(_ grid: [[Int]], _ i: Int, _ j: Int) -> Int {
        if i >= 0 && i < grid.count && j >= 0 && j < grid[0].count {
            return grid[i][j]
        }
        return 0
    }

    func maxSum(_ grid: [[Int]]) -> Int {
        var result = grid[0][0]
        var partialSums = Array(repeating: Array(repeating: 0, count: grid[0].count), count: grid.count)

        for i in 0..<grid.count {
            for j in 0..<grid[0].count {
                partialSums[i][j] = grid[i][j] + get(partialSums, i - 1, j) +
                get(partialSums, i, j - 1) -
                get(partialSums, i - 1, j - 1)
            }
        }

        for x1 in 0..<grid.count {
            for y1 in 0..<grid[0].count {
                for x2 in x1..<grid.count {
                    for y2 in y1..<grid[0].count {
                        let sum = get(partialSums, x2, y2) -
                        get(partialSums, x1 - 1, y2) -
                        get(partialSums, x2, y1 - 1) +
                        get(partialSums, x1 - 1, y1 - 1)
                        result = max(result, sum)
                    }
                }
            }
        }
        return result
    }
}

The Swift details worth knowing:

Arrays of arrays are values. Array(repeating: Array(repeating: 0, count: c), count: r) copies the inner array into every row. Because arrays are value types, changing partialSums[0][0] does not touch the other rows. The result is an independent grid, not r references to one row.
Unlabeled parameters. The _ before each parameter name means calls read get(partialSums, i - 1, j), with no labels.
Calling a method from a method. Inside the class, get(...) means self.get(...), so no prefix is needed.
Lines that continue. A long expression is split here by ending each broken line with a + or -, so the compiler knows the expression goes on in the next line.
Passing the grid in. The arrays are passed by value, but Swift shares the storage until something writes to it, so get does not copy the whole grid on each call.
Empty grids crash. grid[0] on an empty array stops the program with a run-time error. The lesson's tests use grids with at least one cell.

Complexity

Let the grid have m rows and n columns.

Time: O(m² * n²). Building the prefix table is O(m * n). After that, the four corner loops pick about m² * n² / 4 rectangles, and each one costs four constant-time lookups. Without prefix sums, every rectangle sum would cost up to m * n more.

Space: O(m * n). The table partialSums has as many cells as the grid.

Edge cases and common mistakes

All negative numbers. [[-3, -1], [-2, -5]] returns -1, the best single cell. This is why result starts from grid[0][0] and not 0.
One cell. [[7]] returns 7.
A single row or column. [[1, 2, 3]] returns 6, and [[1], [-2], [3]] returns 3.
Forgetting the added corner. Leaving out + get(partialSums, x1 - 1, y1 - 1) subtracts the top-left region twice, so every rectangle that does not touch the top or left edge is off by the sum of that region.
Wrong loop starts. x2 must start at x1 and y2 at y1. Starting them at 0 produces rectangles whose bottom-right corner is above or left of the top-left one, and their "sums" are meaningless.
Out-of-range neighbours. Reading partialSums[i - 1][j] directly would crash on row 0. The helper is what avoids that.

Testing

The grid from the example gives 38.
[[-3, -1], [-2, -5]] gives -1.
[[7]] gives 7.
[[1, 2, 3]] gives 6.
[[1], [-2], [3]] gives 3.

Keep going with SwiftAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.

Start learning free