题解 | #小美的平衡矩阵#

我的思路是把 1 当作 1，0 当作 -1，1 和 0 的个数相等等价于矩阵中所有元素之和为 0。

对于 k 从 1 到 n，遍历所有矩阵，计算和为 0 的矩阵的个数，就是 k * k 的完美矩阵的个数。

现在的问题是：如何快速计算从 (i, j) 开始，大小为 k * k 的矩阵的元素的和? (将其设为 f(i, j, k))

当前要计算的 k * k 矩阵的元素和 = 左上方的 (k - 1) * (k - 1) 矩阵的元素和 + 最下面一行元素的和 + 最右边一列元素的和 - 右下方的元素（因为重复算了一次）

设 matrix(i, j) 表示矩阵第 i 行第 j 列的元素，rowPrefix(i, j) 表示第 i 行前 j 个元素的和，colPrefix(i, j) 表示第 i 列前 j 个元素的和，那么由前缀和有：

f(i, j, k) = f(i, j, k - 1) + rowPrefix(i + k - 1, j + k - 1) - rowPrefix(i + k - 1, j - 1) + colPrefix(j + k - 1, i + k - 1) - colPrefix(j + k - 1, i - 1) - matrix(i, j)

根据此公式Js答案如下：

const rl = require("readline").createInterface({ input: process.stdin });
var iter = rl[Symbol.asyncIterator]();
const readline = async () => (await iter.next()).value;

void async function () {
    // 读取矩阵
    const n = parseInt(await readline())
    const matrix = new Array(n + 1)
    for(let i = 1; i <= n; i++) {
        matrix[i] = new Array(n + 1)
        const nums = (await readline()).split('')
        for(let j = 0; j < nums.length; j++){
            // '1' 存 1，'0' 存 -1，方便算和
            matrix[i][j + 1] = nums[j] === '1' ? 1 : -1
        }
    }

    // 初始化行和列的前缀和
    const rowPrefix = new Array(n + 1)
    const colPrefix = new Array(n + 1)
    for(let i = 1; i <= n; i++) {
        rowPrefix[i] = new Array(n + 1)
        colPrefix[i] = new Array(n + 1)
        rowPrefix[i][0] = 0
        colPrefix[i][0] = 0
        for(let j = 1; j <= n; j++) {
            rowPrefix[i][j] = rowPrefix[i][j - 1] + matrix[i][j]
            colPrefix[i][j] = colPrefix[i][j - 1] + matrix[j][i]
        }
    }

    // 初始化 map 记录前缀和
    let map = new Map()
    for(let i = 1; i <= n; i++) {
        for(let j = 1; j <= n; j++) {
            map.set(`${i} ${j}`, matrix[i][j])
        }
    }
    console.log(0)

    for(let k = 2; k <= n; k++) {
        let newMap = new Map()
        let res, count = 0, row, col
        for(let i = 1; i <= n - k + 1; i++) {
            for(let j = 1; j <= n - k + 1; j++) {
                row = i + k - 1
                col = j + k - 1
                res = map.get(`${i} ${j}`) + rowPrefix[row][col] - rowPrefix[row][j - 1] + colPrefix[col][row] - colPrefix[col][i - 1] - matrix[row][col]
                if(res === 0) {
                    count++    
                }
                newMap.set(`${i} ${j}`, res)
            }
        }
        console.log(count)
        // 更新 map
        map = newMap
    }
}()