小美的平衡矩阵
美团 2024-03-09 春招笔试编程题（编程题第 3 题，25 分）
题目描述
小美拿到了一个 n × n 的矩阵，其中每个元素是 0 或者 1。 小美认为一个矩形区域是完美的，当且仅当该区域内 0 的数量恰好等于 1 的数量。 现在，小美希望你回答：存在多少个 i × i 的完美矩形区域？ 你需要回答 1 ≤ i ≤ n 的所有答案。
输入格式
第一行一个正整数 n，代表矩阵大小。 接下来的 n 行，每行一个长度为 n 的 01 串，用来表示矩阵。
输出格式
输出 n 行，第 i 行输出 i × i 的完美矩形区域的数量。
示例
示例 1
输入：
4
1010
0101
1100
0011
输出：
0
7
0
1
说明： 1×1 不可能平衡（只有一个元素），3×3 也不可能（9 个格子无法平分）， 2×2 的完美子矩阵有 7 个，整个 4×4 矩阵恰好有 8 个 1 和 8 个 0，是唯一一个， 所以第 4 行输出 1。
解题思路
第一步：把「0 和 1 数量相等」换成「1 的个数 = 面积的一半」
矩阵里只有 0 和 1，若区域里有 k 个 1，则有 面积 - k 个 0。
两者相等 ⇔ k = 面积 / 2。
于是只要快速求任意子矩阵里 1 的个数即可 —— 二维前缀和 O(1) 查询。
第二步：奇数边长直接是 0
边长 i 为奇数时，面积 i² 是奇数，i²/2 不是整数，
「1 的个数 = i²/2」不可能成立，所以答案恒为 0，可以直接跳过。
第三步：枚举偶数边长 + 枚举右下角
对每个偶数边长 k，枚举所有 k × k 子矩阵的右下角 (i, j)，
用前缀和判断 1 的个数是否为 k²/2，是就累加。
二维前缀和查询公式：
ones = pre[i][j] - pre[i-k][j] - pre[i][j-k] + pre[i-k][j-k]
第四步：复杂度
共有 O(n) 个偶数边长，每个边长要枚举 O(n²) 个右下角， 总复杂度 O(n³)。本题 n 不大（配合 1 秒时限的规模），可以接受。
完整代码实现
算法可视化 · 逐帧回放这道题的代码189 帧 · 录制输入
import sys
def solve() -> None:
"""二维前缀和，枚举偶数边长的子矩阵"""
data = sys.stdin.read().split()
if not data:
return
n = int(data[0])
grid = data[1:1 + n]
pre = [[0] * (n + 1) for _ in range(n + 1)]
for i in range(n):
row = grid[i]
for j in range(n):
pre[i + 1][j + 1] = (
pre[i][j + 1] + pre[i + 1][j] - pre[i][j]
+ (1 if row[j] == '1' else 0)
)
res = [0] * (n + 1)
for k in range(2, n + 1, 2): # 奇数边长必为 0，跳过
half = k * k // 2
for i in range(k, n + 1):
for j in range(k, n + 1):
ones = (pre[i][j] - pre[i - k][j]
- pre[i][j - k] + pre[i - k][j - k])
if ones == half:
res[k] += 1
sys.stdout.write("\n".join(str(res[i]) for i in range(1, n + 1)) + "\n")
if __name__ == "__main__":
solve()
示例推演
以 4×4 的矩阵为例，看几个有代表性的子矩阵：
| 边长 | 区域 | 1 的个数 | 面积的一半 | 是否完美 |
|---|---|---|---|---|
| 1 | 任意单格 | 0 或 1 | 0.5（非整数） | 否 |
| 2 | 左上角 10/01 | 2 | 2 | 是 |
| 2 | 右上角 10/01 | 2 | 2 | 是 |
| 4 | 整个矩阵 | 8 | 8 | 是 |
按边长统计：边长 1 → 0 个，边长 2 → 7 个，边长 3 → 0 个，边长 4 → 1 个。
与样例输出 0 7 0 1 一致。
复杂度分析
| 解法 | 时间复杂度 | 空间复杂度 | 说明 |
|---|---|---|---|
| 暴力求和 | O(n⁵) | O(n²) | 超时 |
| 二维前缀和 | O(n³) | O(n²) | 推荐 |
易错点总结
1. 把偶数边长漏掉
奇数边长一定为 0，但偶数边长必须枚举，别把判断写反。
2. 前缀和边界
前缀和数组开 (n+1) × (n+1)，下标从 1 起，查询时用 i-k、j-k
才不会越界。
3. 输入是一整串 01
每行是一个长度为 n 的字符串（如 1010），不能按整数读，
否则会读成一个数或者溢出。用 sys.stdin.read().split() 拿到字符串即可。
4. 输出是 n 行
即使某个边长的答案是 0，也要输出那一行（共输出 n 行）。
扩展思考
- 若题目问的是「所有矩形」（不要求正方形），二维前缀和依然够用， 但枚举量会变成 O(n⁴)。
- 「0/1 各占一半」这类条件，本质都是把问题转成「求和」后用前缀和加速。