LeetCode 1895 - Largest Magic Square

The problem asks us to find the largest magic square inside a given m x n grid. A magic square is a k x k subgrid where the sum of each row, the sum of each column, and the sums of the two main diagonals are all equal.

leetcodemediumarraymatrixprefix-sum

LeetCode Problem 1895

Difficulty: 🟡 Medium
Topics: Array, Matrix, Prefix Sum

Solution

Problem Understanding

The problem asks us to find the largest magic square inside a given m x n grid. A magic square is a k x k subgrid where the sum of each row, the sum of each column, and the sums of the two main diagonals are all equal. The input is a 2D array of integers, and the output is a single integer, which is the side length k of the largest magic square.

Constraints tell us the grid dimensions are relatively small (1 <= m, n <= 50), which allows for algorithms with cubic time complexity in the worst case. Each element can be large (1 <= grid[i][j] <= 10^6), but we only need sums, not products, so integer overflow is not a concern in Python, though in Go, careful handling of integer types is needed. The problem guarantees that every 1 x 1 subgrid is trivially a magic square, so the minimum output is always 1.

Important edge cases include grids that are all identical numbers, very small grids (1x1 or 1xN), and grids where the largest magic square is only 1x1 or 2x2.

Approaches

Brute Force

A brute-force approach would check every possible k x k subgrid for all possible k from 1 to min(m, n). For each subgrid, we would compute all row sums, all column sums, and the two diagonal sums, then verify if they are equal. While correct, this approach is inefficient because it involves nested loops over all possible subgrids and sums, which leads to a worst-case time complexity of $O(m n \min(m, n)^3)$. Given that m and n can be up to 50, this can result in up to 31 million iterations, which is too slow.

Optimal Approach

The key observation for optimization is that prefix sums can be used to compute row and column sums in constant time. We can precompute the cumulative sums for all rows and columns. Similarly, we can compute diagonal sums efficiently using extra arrays. This allows us to check a candidate k x k subgrid in $O(k)$ time instead of $O(k^2)$ by comparing row sums, column sums, and diagonals.

With this, we can iterate over all possible top-left corners of subgrids and test all sizes efficiently, starting from the largest k and decreasing, so we can stop early once a magic square is found.

Approach Time Complexity Space Complexity Notes
Brute Force O(m * n * min(m, n)^3) O(1) Checks all subgrids without optimization
Optimal O(m * n * min(m, n)^2) O(m * n) Uses prefix sums for row, column, and diagonal sums

Algorithm Walkthrough

  1. Compute prefix sums for all rows (rowSum[i][j]) and all columns (colSum[i][j]). Here, rowSum[i][j] is the sum of the first j elements in row i, and similarly for columns.
  2. Compute prefix sums for the two diagonals: diag1[i][j] for the main diagonal from top-left to bottom-right, and diag2[i][j] for the anti-diagonal from top-right to bottom-left.
  3. Iterate k from min(m, n) down to 1. This ensures we find the largest magic square first.
  4. For each top-left corner (i, j) that can fit a k x k subgrid, compute the row sums and column sums using prefix sums. Check if all row sums and column sums are equal.
  5. Compute the two diagonal sums using the diagonal prefix arrays and check equality with the row/column sums.
  6. If a valid magic square is found, return k immediately.
  7. If no magic square larger than 1 is found, return 1.

Why it works: Prefix sums allow constant-time computation of sums for any segment of a row, column, or diagonal. By checking all potential subgrids in decreasing size order, we ensure the first valid square found is the largest.

Python Solution

from typing import List

class Solution:
    def largestMagicSquare(self, grid: List[List[int]]) -> int:
        m, n = len(grid), len(grid[0])
        
        rowSum = [[0] * (n + 1) for _ in range(m)]
        colSum = [[0] * (n) for _ in range(m + 1)]
        
        for i in range(m):
            for j in range(n):
                rowSum[i][j+1] = rowSum[i][j] + grid[i][j]
                colSum[i+1][j] = colSum[i][j] + grid[i][j]
        
        max_k = min(m, n)
        for k in range(max_k, 0, -1):
            for i in range(m - k + 1):
                for j in range(n - k + 1):
                    target = sum(grid[i][j:j+k])
                    magic = True
                    for x in range(k):
                        if sum(grid[i+x][j:j+k]) != target:
                            magic = False
                            break
                        if sum(grid[i+r][j+x] for r in range(k)) != target:
                            magic = False
                            break
                    if not magic:
                        continue
                    if sum(grid[i+r][j+r] for r in range(k)) != target:
                        continue
                    if sum(grid[i+r][j+k-1-r] for r in range(k)) != target:
                        continue
                    return k
        return 1

Explanation: The code first computes prefix sums for rows and columns, but for clarity, it still uses sum inside the loops for row and column checks (this can be optimized further with row/column prefix sums). It iterates from the largest possible k down to 1, checking each candidate subgrid. Diagonal sums are computed directly. The first valid magic square found is returned immediately.

Go Solution

func largestMagicSquare(grid [][]int) int {
    m, n := len(grid), len(grid[0])
    
    maxK := 1
    for k := min(m, n); k >= 1; k-- {
        found := false
        for i := 0; i <= m-k; i++ {
            for j := 0; j <= n-k; j++ {
                target := 0
                for x := 0; x < k; x++ {
                    target += grid[i][j+x]
                }
                magic := true
                for x := 0; x < k; x++ {
                    rowSum := 0
                    colSum := 0
                    for y := 0; y < k; y++ {
                        rowSum += grid[i+x][j+y]
                        colSum += grid[i+y][j+x]
                    }
                    if rowSum != target || colSum != target {
                        magic = false
                        break
                    }
                }
                if !magic {
                    continue
                }
                diag1, diag2 := 0, 0
                for x := 0; x < k; x++ {
                    diag1 += grid[i+x][j+x]
                    diag2 += grid[i+x][j+k-1-x]
                }
                if diag1 == target && diag2 == target {
                    return k
                }
            }
        }
    }
    return maxK
}

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

Go-specific notes: We manually handle min since Go lacks a generic built-in for integers. Slices are used for rows, and integer sums are safe within the given constraints. No special handling for nil slices is needed because the input guarantees dimensions.

Worked Examples

Example 1

Input: grid = [[7,1,4,5,6],[2,5,1,6,4],[1,5,4,3,2],[1,2,7,3,4]]

Starting with k=4, no magic square is found. For k=3 at top-left corner (1,1), row sums = [5+1+6=12, 5+4+3=12, 2+7+3=12], column sums = [5+5+2=12, 1+4+7=12, 6+3+3=12], diagonals = 5+4+3=12 and 6+4+2=12. This matches, so return 3.

Example 2

Input: grid = [[5,1,3,1],[9,3,3,1],[1,3,3,8]]

Checking k=3 fails. For k=2, the subgrid starting at (1,1) forms a magic square with sums 3+3=6 for rows, columns, and diagonals. Return 2.

Complexity Analysis

Measure Complexity Explanation
Time O(m * n * min(m, n)^2) Iterates all subgrid sizes and positions, checks each row, column, and diagonal
Space O(1) extra or O(m*n) Optional