注:本题不能二分答案。「每行每列的元素和都相等」是一个非常刁钻的要求,可能中间的某个 $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. 矩阵中最大的三个菱形和
专题训练
见下面数据结构题单的「一、前缀和」。
分类题单
如何科学刷题?
- 滑动窗口与双指针(定长/不定长/单序列/双序列/三指针/分组循环)
- 二分算法(二分答案/最小化最大值/最大化最小值/第K小)
- 单调栈(基础/矩形面积/贡献法/最小字典序)
- 网格图(DFS/BFS/综合应用)
- 位运算(基础/性质/拆位/试填/恒等式/思维)
- 图论算法(DFS/BFS/拓扑排序/基环树/最短路/最小生成树/网络流)
- 动态规划(入门/背包/划分/状态机/区间/状压/数位/数据结构优化/树形/博弈/概率期望)
- 常用数据结构(前缀和/差分/栈/队列/堆/字典树/并查集/树状数组/线段树)
- 数学算法(数论/组合/概率期望/博弈/计算几何/随机算法)
- 贪心与思维(基本贪心策略/反悔/区间/字典序/数学/思维/脑筋急转弯/构造)
- 链表、树与回溯(前后指针/快慢指针/DFS/BFS/直径/LCA)
- 字符串(KMP/Z函数/Manacher/字符串哈希/AC自动机/后缀数组/子序列自动机)
我的题解精选(已分类)
欢迎关注 B站@灵茶山艾府