O(1) 计算每个数(Python/Java/C++/Go)
本题和 3315. 构造最小位运算数组 II 是一样的,请看 我的题解。
本题和 3315. 构造最小位运算数组 II 是一样的,请看 我的题解。
前置知识:【图解】一张图秒懂二维前缀和。
预处理二维前缀和后,暴力的做法是写一个三重循环:
这样做的时间复杂度是 $\mathcal{O}(mn\min(m,n))$。
只需改一个地方,就能让算法的时间复杂度变小:
比如现在 $\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。但测试了一下,并没有明显提升。
ans++ 最多执行 $\min(m,n)$ 次,三重循环的时间复杂度为 $\mathcal{O}(mn + \min(m,n)) = \mathcal{O}(mn)$。本题还有一种做法:枚举对角线,在对角线上做 不定长滑动窗口,把正方形的左上角和右下角看作滑动窗口的左右端点。这也可以做到 $\mathcal{O}(mn)$ 时间。
用这个思路解决如下问题:
欢迎在评论区分享你的代码。
见下面数据结构题单的「§1.6 二维前缀和」。
欢迎关注 B站@灵茶山艾府
注:本题不能二分答案。「每行每列的元素和都相等」是一个非常刁钻的要求,可能中间的某个 $k$ 满足要求,$k$ 大一点或小一点都无法让每行每列的元素和都相等。
从大到小枚举 $k$,判断 $\textit{grid}$ 是否存在一个 $k\times k$ 的子矩阵 $M$,满足如下要求:
这些参与求和的元素,在 $\textit{grid}$ 中都是连续的,我们可以用四种前缀和计算:
为什么这里有一些 $+1$?原理在 前缀和 中讲了,是为了兼容子数组恰好是前缀的情况,此时仍然可以用两个前缀和之差算出子数组和,无需特判。
写个三重循环,依次枚举 $k,i,j$,其中 $k\times k$ 子矩阵的左上角为 $(i-k,j-k)$,右下角为 $(i-1,j-1)$,那么:
代码实现时,可以先求主对角线的元素和、反对角线的元素和,如果二者不相等,则无需枚举 $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
}
}
}
}
从大到小枚举 $k$,判断 $\textit{grid}$ 是否存在一个 $k\times k$ 的子矩阵 $M$,满足如下要求:
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
}
见下面数据结构题单的「一、前缀和」。
欢迎关注 B站@灵茶山艾府