[美团春招笔试题]小美的平衡矩阵(golang实现)
题目
小美拿到了一个n∗nn∗n的矩阵,其中每个元素是 0 或者 1。小美认为一个矩形区域是完美的,当且仅当该区域内 0 的数量恰好等于 1 的数量。现在,小美希望你回答有多少个i∗ii∗i的完美矩形区域。你需要回答1≤i≤n1≤i≤n的所有答案。
思路
使用前缀和
- 使用前缀和(Prefix Sum)来快速计算任意子矩阵中 0 和 1 的数量。
代码实现
package main
import (
"fmt"
)
func main() {
var n int
fmt.Scan(&n)
// 读取矩阵并初始化
matrix := make([][]int, n)
for i := range matrix {
matrix[i] = make([]int, n)
var temp string
fmt.Scan(&temp)
for j := 0; j < n; j++ {
if temp[j] == '0' {
matrix[i][j] = 0
} else {
matrix[i][j] = 1
}
}
}
// 计算前缀和
prefix0 := make([][]int, n+1)
prefix1 := make([][]int, n+1)
for i := 0; i <= n; i++ {
prefix0[i] = make([]int, n+1)
prefix1[i] = make([]int, n+1)
}
for i := 1; i <= n; i++ {
for j := 1; j <= n; j++ {
if matrix[i-1][j-1] == 1{
prefix0[i][j] = prefix0[i-1][j] + prefix0[i][j-1] - prefix0[i-1][j-1]
prefix1[i][j] = prefix1[i-1][j] + prefix1[i][j-1] - prefix1[i-1][j-1] + 1
}else{
prefix0[i][j] = prefix0[i-1][j] + prefix0[i][j-1] - prefix0[i-1][j-1] + 1
prefix1[i][j] = prefix1[i-1][j] + prefix1[i][j-1] - prefix1[i-1][j-1]
}
}
}
// 遍历所有可能的子矩阵大小
for k := 1; k <= n; k++ {
if k%2 != 0 {
fmt.Println(0)
continue
}
ans := 0
// 遍历所有可能的子矩阵
for i := 0; i+k <= n; i++ {
for j := 0; j+k <= n; j++ {
// 使用前缀和快速计算子矩阵中 0 和 1 的数量
count0 := prefix0[i+k][j+k] - prefix0[i][j+k] - prefix0[i+k][j] + prefix0[i][j]
count1 := prefix1[i+k][j+k] - prefix1[i][j+k] - prefix1[i+k][j] + prefix1[i][j]
if count0 == count1 {
ans++
}
}
}
fmt.Println(ans)
}
}
代码解析
(1) 前缀和计算
- 使用 prefix0 和 prefix1 两个二维数组来存储 0 和 1 的前缀和。
- 前缀和的计算公式:go复制
(2) 子矩阵的快速计算
- 使用前缀和快速计算任意子矩阵中 0 和 1 的数量:go复制
(3) 减少重复计算
- 通过前缀和避免了每次重新计算子矩阵的值,减少了时间复杂度。
(4) 代码结构优化
- 将前缀和的计算和子矩阵的遍历分开,提高了代码的可读性。
