阅读视图

发现新文章,点击刷新页面。

无需二分,暴力枚举就是 O(mn)(Python/Java/C++/Go)

前置知识【图解】一张图秒懂二维前缀和

预处理二维前缀和后,暴力的做法是写一个三重循环:

  • 外面两重循环,枚举正方形的左上角 $(i,j)$。
  • 最内层循环,枚举正方形的边长为 $1,2,3,\ldots$ 直到出界或者正方形元素和超过 $\textit{threshold}$ 为止。在此过程中更新答案 $\textit{ans}$ 的最大值。

这样做的时间复杂度是 $\mathcal{O}(mn\min(m,n))$。

只需改一个地方,就能让算法的时间复杂度变小:

  • 最内层循环,从 $\textit{ans}+1$ 开始枚举。

比如现在 $\textit{ans}=3$,那么枚举正方形的边长为 $1,2,3$ 是毫无意义的,不会让答案变得更大。所以直接从 $\textit{ans}+1=4$ 开始枚举更好。

###py

class Solution:
    def maxSideLength(self, mat: List[List[int]], threshold: int) -> int:
        m, n = len(mat), len(mat[0])
        s = [[0] * (n + 1) for _ in range(m + 1)]
        for i, row in enumerate(mat):
            for j, x in enumerate(row):
                s[i + 1][j + 1] = s[i + 1][j] + s[i][j + 1] - s[i][j] + x

        # 返回左上角在 (r1, c1),右下角在 (r2, c2) 的子矩阵元素和
        def query(r1: int, c1: int, r2: int, c2: int) -> int:
            return s[r2 + 1][c2 + 1] - s[r2 + 1][c1] - s[r1][c2 + 1] + s[r1][c1]

        ans = 0
        for i in range(m):
            for j in range(n):
                # 边长为 ans+1 的正方形,左上角在 (i, j),右下角在 (i+ans, j+ans)
                while i + ans < m and j + ans < n and query(i, j, i + ans, j + ans) <= threshold:
                    ans += 1
        return ans

###java

class Solution {
    public int maxSideLength(int[][] mat, int threshold) {
        int m = mat.length;
        int n = mat[0].length;
        int[][] sum = new int[m + 1][n + 1];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                sum[i + 1][j + 1] = sum[i + 1][j] + sum[i][j + 1] - sum[i][j] + mat[i][j];
            }
        }

        int ans = 0;
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                // 边长为 ans+1 的正方形,左上角在 (i, j),右下角在 (i+ans, j+ans)
                while (i + ans < m && j + ans < n && query(sum, i, j, i + ans, j + ans) <= threshold) {
                    ans++;
                }
            }
        }
        return ans;
    }

    // 返回左上角在 (r1, c1),右下角在 (r2, c2) 的子矩阵元素和
    private int query(int[][] sum, int r1, int c1, int r2, int c2) {
        return sum[r2 + 1][c2 + 1] - sum[r2 + 1][c1] - sum[r1][c2 + 1] + sum[r1][c1];
    }
}

###cpp

class Solution {
public:
    int maxSideLength(vector<vector<int>>& mat, int threshold) {
        int m = mat.size(), n = mat[0].size();
        vector sum(m + 1, vector<int>(n + 1));
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                sum[i + 1][j + 1] = sum[i + 1][j] + sum[i][j + 1] - sum[i][j] + mat[i][j];
            }
        }

        // 返回左上角在 (r1, c1),右下角在 (r2, c2) 的子矩阵元素和
        auto query = [&](int r1, int c1, int r2, int c2) -> int {
            return sum[r2 + 1][c2 + 1] - sum[r2 + 1][c1] - sum[r1][c2 + 1] + sum[r1][c1];
        };

        int ans = 0;
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                // 边长为 ans+1 的正方形,左上角在 (i, j),右下角在 (i+ans, j+ans)
                while (i + ans < m && j + ans < n && query(i, j, i + ans, j + ans) <= threshold) {
                    ans++;
                }
            }
        }
        return ans;
    }
};

###go

func maxSideLength(mat [][]int, threshold int) (ans int) {
m, n := len(mat), len(mat[0])
sum := make([][]int, m+1)
sum[0] = make([]int, n+1)
for i, row := range mat {
sum[i+1] = make([]int, n+1)
for j, x := range row {
sum[i+1][j+1] = sum[i+1][j] + sum[i][j+1] - sum[i][j] + x
}
}

// 返回左上角在 (r1, c1),右下角在 (r2, c2) 的子矩阵元素和
query := func(r1, c1, r2, c2 int) int {
return sum[r2+1][c2+1] - sum[r2+1][c1] - sum[r1][c2+1] + sum[r1][c1]
}

for i := range m {
for j := range n {
// 边长为 ans+1 的正方形,左上角在 (i, j),右下角在 (i+ans, j+ans)
for i+ans < m && j+ans < n && query(i, j, i+ans, j+ans) <= threshold {
ans++
}
}
}
return
}

注:外层循环可以改成 i + ans < m 以及 j + ans < n。但测试了一下,并没有明显提升。

复杂度分析

  • 时间复杂度:$\mathcal{O}(mn)$,其中 $m$ 和 $n$ 分别是 $\textit{grid}$ 的行数和列数。虽然我们写了个三重循环,但由于答案最大是 $\min(m,n)$,所以最内层的 ans++ 最多执行 $\min(m,n)$ 次,三重循环的时间复杂度为 $\mathcal{O}(mn + \min(m,n)) = \mathcal{O}(mn)$。
  • 空间复杂度:$\mathcal{O}(mn)$。

思考题

本题还有一种做法:枚举对角线,在对角线上做 不定长滑动窗口,把正方形的左上角和右下角看作滑动窗口的左右端点。这也可以做到 $\mathcal{O}(mn)$ 时间。

用这个思路解决如下问题:

  • 计算元素总和小于或等于阈值的正方形的个数

欢迎在评论区分享你的代码。

专题训练

见下面数据结构题单的「§1.6 二维前缀和」。

分类题单

如何科学刷题?

  1. 滑动窗口与双指针(定长/不定长/单序列/双序列/三指针/分组循环)
  2. 二分算法(二分答案/最小化最大值/最大化最小值/第K小)
  3. 单调栈(基础/矩形面积/贡献法/最小字典序)
  4. 网格图(DFS/BFS/综合应用)
  5. 位运算(基础/性质/拆位/试填/恒等式/思维)
  6. 图论算法(DFS/BFS/拓扑排序/基环树/最短路/最小生成树/网络流)
  7. 动态规划(入门/背包/划分/状态机/区间/状压/数位/数据结构优化/树形/博弈/概率期望)
  8. 常用数据结构(前缀和/差分/栈/队列/堆/字典树/并查集/树状数组/线段树)
  9. 数学算法(数论/组合/概率期望/博弈/计算几何/随机算法)
  10. 贪心与思维(基本贪心策略/反悔/区间/字典序/数学/思维/脑筋急转弯/构造)
  11. 链表、树与回溯(前后指针/快慢指针/DFS/BFS/直径/LCA)
  12. 字符串(KMP/Z函数/Manacher/字符串哈希/AC自动机/后缀数组/子序列自动机)

我的题解精选(已分类)

欢迎关注 B站@灵茶山艾府

从 O(N^4) 优化到 O(N^3)(Python/Java/C++/Go)

注:本题不能二分答案。「每行每列的元素和都相等」是一个非常刁钻的要求,可能中间的某个 $k$ 满足要求,$k$ 大一点或小一点都无法让每行每列的元素和都相等。

方法一:四种前缀和

从大到小枚举 $k$,判断 $\textit{grid}$ 是否存在一个 $k\times k$ 的子矩阵 $M$,满足如下要求:

  • 设 $M$ 第一行的元素和为 $s$。
  • $M$ 每行的元素和都是 $s$。
  • $M$ 每列的元素和都是 $s$。
  • $M$ 主对角线的元素和为 $s$。
  • $M$ 反对角线的元素和为 $s$。

这些参与求和的元素,在 $\textit{grid}$ 中都是连续的,我们可以用四种前缀和计算:

  • $\textit{rowSum}[i][j+1]$ 表示 $\textit{grid}$ 的 $i$ 行的前缀 $[0,j]$ 的元素和,即 $(i,0),(i,1),\ldots,(i,j)$ 的元素和。
  • $\textit{colSum}[i+1][j]$ 表示 $\textit{grid}$ 的 $j$ 列的前缀 $[0,i]$ 的元素和,即 $(0,j),(1,j),\ldots,(i,j)$ 的元素和。
  • $\textit{diagSum}[i+1][j+1]$ 表示从最上边或最左边出发,向右下↘到 $(i,j)$ 这条线上的元素和。
  • $\textit{antiSum}[i+1][j]$ 表示从最上边或最右边出发,向左下↙到 $(i,j)$ 这条线上的元素和。

为什么这里有一些 $+1$?原理在 前缀和 中讲了,是为了兼容子数组恰好是前缀的情况,此时仍然可以用两个前缀和之差算出子数组和,无需特判。

写个三重循环,依次枚举 $k,i,j$,其中 $k\times k$ 子矩阵的左上角为 $(i-k,j-k)$,右下角为 $(i-1,j-1)$,那么:

  • 主对角线的元素和为 $\textit{diagSum}[i][j] - \textit{diagSum}[i-k][j-k]$。
  • 反对角线的元素和为 $\textit{antiSum}[i][j-k]-\textit{antiSum}[i-k][j]$。
  • 在 $[i-k,i-1]$ 中枚举行号 $r$,行元素和为 $\textit{rowSum}[r][j] - \textit{rowSum}[r][j-k]$。
  • 在 $[j-k,j-1]$ 中枚举列号 $c$,列元素和为 $\textit{colSum}[i][c] - \textit{colSum}[i-k][c]$。

代码实现时,可以先求主对角线的元素和、反对角线的元素和,如果二者不相等,则无需枚举 $r$ 和 $c$。

class Solution:
    def largestMagicSquare(self, grid: List[List[int]]) -> int:
        m, n = len(grid), len(grid[0])
        row_sum = [[0] * (n + 1) for _ in range(m)]       # → 前缀和
        col_sum = [[0] * n for _ in range(m + 1)]         # ↓ 前缀和
        diag_sum = [[0] * (n + 1) for _ in range(m + 1)]  # ↘ 前缀和
        anti_sum = [[0] * (n + 1) for _ in range(m + 1)]  # ↙ 前缀和

        for i, row in enumerate(grid):
            for j, x in enumerate(row):
                row_sum[i][j + 1] = row_sum[i][j] + x
                col_sum[i + 1][j] = col_sum[i][j] + x
                diag_sum[i + 1][j + 1] = diag_sum[i][j] + x
                anti_sum[i + 1][j] = anti_sum[i][j + 1] + x

        # k×k 子矩阵的左上角为 (i−k, j−k),右下角为 (i−1, j−1)
        for k in range(min(m, n), 0, -1):
            for i in range(k, m + 1):
                for j in range(k, n + 1):
                    # 子矩阵主对角线的和
                    s = diag_sum[i][j] - diag_sum[i - k][j - k]

                    # 子矩阵反对角线的和等于 s
                    # 子矩阵每行的和都等于 s
                    # 子矩阵每列的和都等于 s
                    if anti_sum[i][j - k] - anti_sum[i - k][j] == s and \
                       all(row_sum[r][j] - row_sum[r][j - k] == s for r in range(i - k, i)) and \
                       all(col_sum[i][c] - col_sum[i - k][c] == s for c in range(j - k, j)):
                        return k
class Solution {
    public int largestMagicSquare(int[][] grid) {
        int m = grid.length;
        int n = grid[0].length;
        int[][] rowSum = new int[m][n + 1];      // → 前缀和
        int[][] colSum = new int[m + 1][n];      // ↓ 前缀和
        int[][] diagSum = new int[m + 1][n + 1]; // ↘ 前缀和
        int[][] antiSum = new int[m + 1][n + 1]; // ↙ 前缀和

        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                int x = grid[i][j];
                rowSum[i][j + 1] = rowSum[i][j] + x;
                colSum[i + 1][j] = colSum[i][j] + x;
                diagSum[i + 1][j + 1] = diagSum[i][j] + x;
                antiSum[i + 1][j] = antiSum[i][j + 1] + x;
            }
        }

        // k×k 子矩阵的左上角为 (i−k, j−k),右下角为 (i−1, j−1)
        for (int k = Math.min(m, n); ; k--) {
            for (int i = k; i <= m; i++) {
                next:
                for (int j = k; j <= n; j++) {
                    // 子矩阵主对角线的和
                    int sum = diagSum[i][j] - diagSum[i - k][j - k];

                    // 子矩阵反对角线的和
                    if (antiSum[i][j - k] - antiSum[i - k][j] != sum) {
                        continue;
                    }

                    // 子矩阵每行的和
                    for (int r = i - k; r < i; r++) {
                        if (rowSum[r][j] - rowSum[r][j - k] != sum) {
                            continue next;
                        }
                    }

                    // 子矩阵每列的和
                    for (int c = j - k; c < j; c++) {
                        if (colSum[i][c] - colSum[i - k][c] != sum) {
                            continue next;
                        }
                    }

                    return k;
                }
            }
        }
    }
}
class Solution {
public:
    int largestMagicSquare(vector<vector<int>>& grid) {
        int m = grid.size(), n = grid[0].size();
        vector row_sum(m, vector<int>(n + 1));      // → 前缀和
        vector col_sum(m + 1, vector<int>(n));      // ↓ 前缀和
        vector diag_sum(m + 1, vector<int>(n + 1)); // ↘ 前缀和
        vector anti_sum(m + 1, vector<int>(n + 1)); // ↙ 前缀和

        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                int x = grid[i][j];
                row_sum[i][j + 1] = row_sum[i][j] + x;
                col_sum[i + 1][j] = col_sum[i][j] + x;
                diag_sum[i + 1][j + 1] = diag_sum[i][j] + x;
                anti_sum[i + 1][j] = anti_sum[i][j + 1] + x;
            }
        }

        // k×k 子矩阵的左上角为 (i−k, j−k),右下角为 (i−1, j−1)
        for (int k = min(m, n); ; k--) {
            for (int i = k; i <= m; i++) {
                for (int j = k; j <= n; j++) {
                    // 子矩阵主对角线的和
                    int sum = diag_sum[i][j] - diag_sum[i - k][j - k];

                    // 子矩阵反对角线的和
                    if (anti_sum[i][j - k] - anti_sum[i - k][j] != sum) {
                        continue;
                    }

                    // 子矩阵每行的和
                    bool ok = true;
                    for (int r = i - k; r < i; r++) {
                        if (row_sum[r][j] - row_sum[r][j - k] != sum) {
                            ok = false;
                            break;
                        }
                    }
                    if (!ok) {
                        continue;
                    }

                    // 子矩阵每列的和
                    for (int c = j - k; c < j; c++) {
                        if (col_sum[i][c] - col_sum[i - k][c] != sum) {
                            ok = false;
                            break;
                        }
                    }
                    if (ok) {
                        return k;
                    }
                }
            }
        }
    }
};
func largestMagicSquare(grid [][]int) int {
m, n := len(grid), len(grid[0])
rowSum := make([][]int, m)    // → 前缀和
colSum := make([][]int, m+1)  // ↓ 前缀和
diagSum := make([][]int, m+1) // ↘ 前缀和
antiSum := make([][]int, m+1) // ↙ 前缀和
for i := range m + 1 {
colSum[i] = make([]int, n)
diagSum[i] = make([]int, n+1)
antiSum[i] = make([]int, n+1)
}

for i, row := range grid {
rowSum[i] = make([]int, n+1)
for j, x := range row {
rowSum[i][j+1] = rowSum[i][j] + x
colSum[i+1][j] = colSum[i][j] + x
diagSum[i+1][j+1] = diagSum[i][j] + x
antiSum[i+1][j] = antiSum[i][j+1] + x
}
}

// k×k 子矩阵的左上角为 (i−k, j−k),右下角为 (i−1, j−1)
for k := min(m, n); ; k-- {
for i := k; i <= m; i++ {
next:
for j := k; j <= n; j++ {
// 子矩阵主对角线的和
sum := diagSum[i][j] - diagSum[i-k][j-k]

// 子矩阵反对角线的和
if antiSum[i][j-k]-antiSum[i-k][j] != sum {
continue
}

// 子矩阵每行的和
for _, rowS := range rowSum[i-k : i] {
if rowS[j]-rowS[j-k] != sum {
continue next
}
}

// 子矩阵每列的和
for c := j - k; c < j; c++ {
if colSum[i][c]-colSum[i-k][c] != sum {
continue next
}
}

return k
}
}
}
}

复杂度分析

  • 时间复杂度:$\mathcal{O}(mn\min(m,n)^2)$,其中 $m$ 和 $n$ 分别是 $\textit{grid}$ 的行数和列数。
  • 空间复杂度:$\mathcal{O}(mn)$。

方法二:维护连续等和行列的个数

从大到小枚举 $k$,判断 $\textit{grid}$ 是否存在一个 $k\times k$ 的子矩阵 $M$,满足如下要求:

  • 设 $M$ 第一行的元素和为 $s$。
  • $M$ 每行的元素和都是 $s$。优化:想象有一个 $k\times k$ 的窗口在向下滑动,我们可以维护到第 $i$ 行时,有连续多少行的和都等于 $s$。维护一个计数器 $\textit{sameCnt}$,如果当前行的和等于前一行的和,那么把 $\textit{sameCnt}$ 加一,否则把 $\textit{sameCnt}$ 重置为 $1$。如果 $\textit{sameCnt}\ge k$,则说明子矩阵每行的元素和都相等。
  • $M$ 每列的元素和都是 $s$。优化:想象有一个 $k\times k$ 的窗口在向右滑动,我们可以维护到第 $j$ 列时,有连续多少列的和都等于 $s$。算法同上。
  • $M$ 主对角线的元素和为 $s$。
  • $M$ 反对角线的元素和为 $s$。
class Solution:
    def largestMagicSquare(self, grid: List[List[int]]) -> int:
        m, n = len(grid), len(grid[0])
        row_sum = [[0] * (n + 1) for _ in range(m)]       # → 前缀和
        col_sum = [[0] * n for _ in range(m + 1)]         # ↓ 前缀和
        diag_sum = [[0] * (n + 1) for _ in range(m + 1)]  # ↘ 前缀和
        anti_sum = [[0] * (n + 1) for _ in range(m + 1)]  # ↙ 前缀和

        for i, row in enumerate(grid):
            for j, x in enumerate(row):
                row_sum[i][j + 1] = row_sum[i][j] + x
                col_sum[i + 1][j] = col_sum[i][j] + x
                diag_sum[i + 1][j + 1] = diag_sum[i][j] + x
                anti_sum[i + 1][j] = anti_sum[i][j + 1] + x

        # is_same_col_sum[i][j] 表示右下角为 (i, j) 的子矩形,每列元素和是否都相等
        is_same_col_sum = [[False] * n for _ in range(m)]

        for k in range(min(m, n), 1, -1):
            for i in range(k, m + 1):
                # 想象有一个 k×k 的窗口在向右滑动
                same_cnt = 1
                for j in range(1, n):
                    if col_sum[i][j] - col_sum[i - k][j] == col_sum[i][j - 1] - col_sum[i - k][j - 1]:
                        same_cnt += 1
                    else:
                        same_cnt = 1
                    # 连续 k 列元素和是否都一样
                    is_same_col_sum[i - 1][j] = same_cnt >= k

            for j in range(k, n + 1):
                # 想象有一个 k×k 的窗口在向下滑动
                sum_row = row_sum[0][j] - row_sum[0][j - k]
                same_cnt = 1
                for i in range(2, m + 1):
                    row_s = row_sum[i - 1][j] - row_sum[i - 1][j - k]
                    if row_s == sum_row:
                        same_cnt += 1
                        if (same_cnt >= k and  # 连续 k 行元素和都一样
                            is_same_col_sum[i - 1][j - 1] and  # 连续 k 列元素和都一样
                            col_sum[i][j - 1] - col_sum[i - k][j - 1] == sum_row and  # 列和 = 行和
                            diag_sum[i][j] - diag_sum[i - k][j - k] == sum_row and  # 主对角线和 = 行和
                            anti_sum[i][j - k] - anti_sum[i - k][j] == sum_row):  # 反对角线和 = 行和
                            return k
                    else:
                        sum_row = row_s
                        same_cnt = 1

        return 1
class Solution {
    public int largestMagicSquare(int[][] grid) {
        int m = grid.length;
        int n = grid[0].length;
        int[][] rowSum = new int[m][n + 1];      // → 前缀和
        int[][] colSum = new int[m + 1][n];      // ↓ 前缀和
        int[][] diagSum = new int[m + 1][n + 1]; // ↘ 前缀和
        int[][] antiSum = new int[m + 1][n + 1]; // ↙ 前缀和

        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                int x = grid[i][j];
                rowSum[i][j + 1] = rowSum[i][j] + x;
                colSum[i + 1][j] = colSum[i][j] + x;
                diagSum[i + 1][j + 1] = diagSum[i][j] + x;
                antiSum[i + 1][j] = antiSum[i][j + 1] + x;
            }
        }

        // isSameColSum[i][j] 表示右下角为 (i, j) 的子矩形,每列元素和是否都相等
        boolean[][] isSameColSum = new boolean[m][n];

        for (int k = Math.min(m, n); k > 1; k--) {
            for (int i = k; i <= m; i++) {
                // 想象有一个 k×k 的窗口在向右滑动
                int sameCnt = 1;
                for (int j = 1; j < n; j++) {
                    if (colSum[i][j] - colSum[i - k][j] == colSum[i][j - 1] - colSum[i - k][j - 1]) {
                        sameCnt++;
                    } else {
                        sameCnt = 1;
                    }
                    // 连续 k 列元素和是否都一样
                    isSameColSum[i - 1][j] = sameCnt >= k;
                }
            }

            for (int j = k; j <= n; j++) {
                // 想象有一个 k×k 的窗口在向下滑动
                int sum = rowSum[0][j] - rowSum[0][j - k];
                int sameCnt = 1;
                for (int i = 2; i <= m; i++) {
                    int rowS = rowSum[i - 1][j] - rowSum[i - 1][j - k];
                    if (rowS == sum) {
                        sameCnt++;
                        if (sameCnt >= k && // 连续 k 行元素和都一样
                            isSameColSum[i - 1][j - 1] && // 连续 k 列元素和都一样
                            colSum[i][j - 1] - colSum[i - k][j - 1] == sum && // 列和 = 行和
                            diagSum[i][j] - diagSum[i - k][j - k] == sum && // 主对角线和 = 行和
                            antiSum[i][j - k] - antiSum[i - k][j] == sum) { // 反对角线和 = 行和
                            return k;
                        }
                    } else {
                        sum = rowS;
                        sameCnt = 1;
                    }
                }
            }
        }

        return 1;
    }
}
class Solution {
public:
    int largestMagicSquare(vector<vector<int>>& grid) {
        int m = grid.size(), n = grid[0].size();
        vector row_sum(m, vector<int>(n + 1));      // → 前缀和
        vector col_sum(m + 1, vector<int>(n));      // ↓ 前缀和
        vector diag_sum(m + 1, vector<int>(n + 1)); // ↘ 前缀和
        vector anti_sum(m + 1, vector<int>(n + 1)); // ↙ 前缀和

        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                int x = grid[i][j];
                row_sum[i][j + 1] = row_sum[i][j] + x;
                col_sum[i + 1][j] = col_sum[i][j] + x;
                diag_sum[i + 1][j + 1] = diag_sum[i][j] + x;
                anti_sum[i + 1][j] = anti_sum[i][j + 1] + x;
            }
        }

        // is_same_col_sum[i][j] 表示右下角为 (i, j) 的子矩形,每列元素和是否都相等
        vector is_same_col_sum(m, vector<int8_t>(n));

        for (int k = min(m, n); k > 1; k--) {
            for (int i = k; i <= m; i++) {
                // 想象有一个 k×k 的窗口在向右滑动
                int same_cnt = 1;
                for (int j = 1; j < n; j++) {
                    if (col_sum[i][j] - col_sum[i - k][j] == col_sum[i][j - 1] - col_sum[i - k][j - 1]) {
                        same_cnt++;
                    } else {
                        same_cnt = 1;
                    }
                    // 连续 k 列元素和是否都一样
                    is_same_col_sum[i - 1][j] = same_cnt >= k;
                }
            }

            for (int j = k; j <= n; j++) {
                // 想象有一个 k×k 的窗口在向下滑动
                int sum_row = row_sum[0][j] - row_sum[0][j - k];
                int same_cnt = 1;
                for (int i = 2; i <= m; i++) {
                    int row_s = row_sum[i - 1][j] - row_sum[i - 1][j - k];
                    if (row_s == sum_row) {
                        same_cnt++;
                        if (same_cnt >= k && // 连续 k 行元素和都一样
                            is_same_col_sum[i - 1][j - 1] && // 连续 k 列元素和都一样
                            col_sum[i][j - 1] - col_sum[i - k][j - 1] == sum_row && // 列和 = 行和
                            diag_sum[i][j] - diag_sum[i - k][j - k] == sum_row && // 主对角线和 = 行和
                            anti_sum[i][j - k] - anti_sum[i - k][j] == sum_row) { // 反对角线和 = 行和
                            return k;
                        }
                    } else {
                        sum_row = row_s;
                        same_cnt = 1;
                    }
                }
            }
        }

        return 1;
    }
};
func largestMagicSquare(grid [][]int) int {
m, n := len(grid), len(grid[0])
rowSum := make([][]int, m)    // → 前缀和
colSum := make([][]int, m+1)  // ↓ 前缀和
diagSum := make([][]int, m+1) // ↘ 前缀和
antiSum := make([][]int, m+1) // ↙ 前缀和
for i := range m + 1 {
colSum[i] = make([]int, n)
diagSum[i] = make([]int, n+1)
antiSum[i] = make([]int, n+1)
}
for i, row := range grid {
rowSum[i] = make([]int, n+1)
for j, x := range row {
rowSum[i][j+1] = rowSum[i][j] + x
colSum[i+1][j] = colSum[i][j] + x
diagSum[i+1][j+1] = diagSum[i][j] + x
antiSum[i+1][j] = antiSum[i][j+1] + x
}
}

// isSameColSum[i][j] 表示右下角为 (i, j) 的子矩形,每列元素和是否都相等
isSameColSum := make([][]bool, m)
for i := range isSameColSum {
isSameColSum[i] = make([]bool, n)
}
for k := min(m, n); k > 1; k-- {
for i := k; i <= m; i++ {
// 想象有一个 k×k 的窗口在向右滑动
sameCnt := 1
for j := 1; j < n; j++ {
if colSum[i][j]-colSum[i-k][j] == colSum[i][j-1]-colSum[i-k][j-1] {
sameCnt++
} else {
sameCnt = 1
}
// 连续 k 列元素和是否都一样
isSameColSum[i-1][j] = sameCnt >= k
}
}

for j := k; j <= n; j++ {
// 想象有一个 k×k 的窗口在向下滑动
sum := rowSum[0][j] - rowSum[0][j-k]
sameCnt := 1
for i := 2; i <= m; i++ {
rowS := rowSum[i-1][j] - rowSum[i-1][j-k]
if rowS == sum {
sameCnt++
if sameCnt >= k && // 连续 k 行元素和都一样
isSameColSum[i-1][j-1] && // 连续 k 列元素和都一样
colSum[i][j-1]-colSum[i-k][j-1] == sum && // 列和 = 行和
diagSum[i][j]-diagSum[i-k][j-k] == sum && // 主对角线和 = 行和
antiSum[i][j-k]-antiSum[i-k][j] == sum {  // 反对角线和 = 行和
return k
}
} else {
sum = rowS
sameCnt = 1
}
}
}
}

return 1
}

复杂度分析

  • 时间复杂度:$\mathcal{O}(mn\min(m,n))$,其中 $m$ 和 $n$ 分别是 $\textit{grid}$ 的行数和列数。
  • 空间复杂度:$\mathcal{O}(mn)$。

相似题目

1878. 矩阵中最大的三个菱形和

专题训练

见下面数据结构题单的「一、前缀和」。

分类题单

如何科学刷题?

  1. 滑动窗口与双指针(定长/不定长/单序列/双序列/三指针/分组循环)
  2. 二分算法(二分答案/最小化最大值/最大化最小值/第K小)
  3. 单调栈(基础/矩形面积/贡献法/最小字典序)
  4. 网格图(DFS/BFS/综合应用)
  5. 位运算(基础/性质/拆位/试填/恒等式/思维)
  6. 图论算法(DFS/BFS/拓扑排序/基环树/最短路/最小生成树/网络流)
  7. 动态规划(入门/背包/划分/状态机/区间/状压/数位/数据结构优化/树形/博弈/概率期望)
  8. 常用数据结构(前缀和/差分/栈/队列/堆/字典树/并查集/树状数组/线段树)
  9. 数学算法(数论/组合/概率期望/博弈/计算几何/随机算法)
  10. 贪心与思维(基本贪心策略/反悔/区间/字典序/数学/思维/脑筋急转弯/构造)
  11. 链表、树与回溯(前后指针/快慢指针/DFS/BFS/直径/LCA)
  12. 字符串(KMP/Z函数/Manacher/字符串哈希/AC自动机/后缀数组/子序列自动机)

我的题解精选(已分类)

欢迎关注 B站@灵茶山艾府

❌