[美团春招笔试题]小美的平衡矩阵(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) 代码结构优化
- 将前缀和的计算和子矩阵的遍历分开,提高了代码的可读性。