[美团春招笔试题]小美的平衡矩阵(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) 代码结构优化

  • 将前缀和的计算和子矩阵的遍历分开,提高了代码的可读性。
#美团#
全部评论

相关推荐

评论
点赞
1
分享

创作者周榜

更多
正在热议
更多
牛客网
牛客企业服务