Largest Magic Square LeetCode solution in Python

1
# Runtime: 1987 ms (Top 15.7%) | Memory: 16.60 MB (Top 82.4%)
2

3

4
class Solution:
5
    def largestMagicSquare(self, grid: List[List[int]]) -> int:
6
        m, n = len(grid), len(grid[0])  # dimensions
7
        rows = [[0] * (n + 1) for _ in range(m)]  # prefix sum along row
8
        cols = [[0] * n for _ in range(m + 1)]  # prefix sum along column
9

10
        for i in range(m):
11
            for j in range(n):
12
                rows[i][j + 1] = grid[i][j] + rows[i][j]
13
                cols[i + 1][j] = grid[i][j] + cols[i][j]
14

15
        ans = 1
16
        for i in range(m):
17
            for j in range(n):
18
                diag = grid[i][j]
19
                for k in range(min(i, j)):
20
                    ii, jj = i - k - 1, j - k - 1
21
                    diag += grid[ii][jj]
22
                    ss = {diag}
23
                    for r in range(ii, i + 1):
24
                        ss.add(rows[r][j + 1] - rows[r][jj])
25
                    for c in range(jj, j + 1):
26
                        ss.add(cols[i + 1][c] - cols[ii][c])
27
                    ss.add(
28
                        sum(grid[ii + kk][j - kk] for kk in range(k + 2))
29
                    )  # anti-diagonal
30
                    if len(ss) == 1:
31
                        ans = max(ans, k + 2)
32
        return ans