###python3

class Solution:
    def largestMagicSquare(self, grid: List[List[int]]) -> int:
        N = len(grid)
        M = len(grid[0])
        up = [[0] * M for _ in range(N)]
        left = [[0] * M for _ in range(N)]
        up_left = [[0] * M for _ in range(N)]
        up_right = [[0] * M for _ in range(N)]

        def check_get(arr, i, j):
            if i >= 0 and 0 <= j < M:
                return arr[i][j]
            else:
                return 0

        for i in range(N):
            for j in range(M):
                up[i][j] = check_get(up, i - 1, j) + grid[i][j]
                left[i][j] = check_get(left, i, j - 1) + grid[i][j]
                up_left[i][j] = check_get(up_left, i - 1, j - 1) + grid[i][j]
                up_right[i][j] = check_get(up_right, i - 1, j + 1) + grid[i][j]
        for k in range(min(M, N), 1, -1):
            candidates = set()
            for i in range(k - 1, N):
                last = up[i][0] - check_get(up, i - k, 0)
                count = 1
                for j in range(1, M):
                    curr = up[i][j] - check_get(up, i - k, j)
                    if curr == last:
                        count += 1
                    else:
                        last = curr
                        count = 1
                    if count >= k:
                        # Check diagonal
                        if up_left[i][j] - check_get(up_left, i - k, j - k) == last\
                                and up_right[i][j - k + 1] - check_get(up_right, i - k, j + 1) == last:
                            candidates.add((i, j))
            if candidates:
                for j in range(k - 1, M):
                    last = left[0][j] - check_get(left, 0, j - k)
                    count = 1
                    for i in range(1, N):
                        curr = left[i][j] - check_get(left, i, j - k)
                        if curr == last:
                            count += 1
                        else:
                            last = curr
                            count = 1
                        if count >= k and (i, j) in candidates:
                            return k
        else:
            return 1

算法原理并不算复杂，仍然是用前缀和来优化计算，不过这里有个额外的优化：

选定某个k的情况下，如果某个k*k的正方形中每一列的和都相等，我们称之为列准幻方。现在希望一次找到所有的列准幻方，首先穷举幻方的最后一行（由于k已经选定第一行也就确定了），然后扫描每一列，注意到新扫描进来的这一列的和可以用前缀和在O(1)时间内计算出来，同时可以立即判断出它是否和前一列相等，随时记录当前相等的列数，就可以判断出以当前列为最后一列的k*k正方形是不是一个列准幻方。对于找到的每个列准幻方，接下来可以校验它的两条对角线是不是和最后一列的和相等，通过前缀和也可以做到O(1)复杂度，这样可以筛选出所有对角线也符合条件的列准幻方。

如果对于每个列准幻方都校验各行的和，则复杂度会变成四次方级别。这里有个非常简单的技巧解决这个问题：我们把行列倒换，用相同的方法求出所有的行准幻方，然后用hash表判断每个行准幻方是否同时也是列准幻方。显然如果一个正方形既是行准幻方又是列准幻方，同时对角线也符合条件，那么它就是一个幻方。

这样总的复杂度就可以降到O(NM min(M,N))。