矩阵计数

[题目链接](https://www.nowcoder.com/practice/3f97500b339d46d7ad969ae414cf87ad)

思路

给定一个 $n \times n$ 的字符矩阵，每个元素是 r、e、d 之一。定义矩阵权值为矩阵中三种字符出现次数的最小值。给定阈值 $myval$ ，求有多少个子方阵的权值不小于 $myval$ 。

二维前缀和 + 二分查找

核心观察：固定子方阵的左上角 $(i, j)$ ，随着边长 $k$ 的增大，方阵内每种字符的数量单调不减，因此权值（三种字符数量的最小值）也单调不减。这意味着可以对边长进行二分查找，找到使权值 $\geq myval$ 的最小边长 $k_0$ ，那么从 $k_0$ 到最大可能边长的所有方阵都满足条件。

快速查询子矩阵字符数：对 r、e、d 三种字符分别建立二维前缀和。设 $pre[c][i][j]$ 为矩阵前 $i$ 行、前 $j$ 列中字符 $c$ 的出现次数，则左上角 $(r_1, c_1)$ 、右下角 $(r_2, c_2)$ 的子矩阵中字符 $c$ 的数量为：

$ $pre[c][r_2][c_2] - pre[c][r_1{-}1][c_2] - pre[c][r_2][c_1{-}1] + pre[c][r_1{-}1][c_1{-}1]$ $

算法流程

构建三种字符的二维前缀和数组。
枚举每个位置 $(i, j)$ 作为子方阵的左上角。
二分查找满足条件的最小边长 $k_0$ ：在 $[1, \min(n-i+1, n-j+1)]$ 上二分，用前缀和 $O(1)$ 查询子方阵中每种字符的数量，取最小值与 $myval$ 比较。
若 $k_0$ 存在，则贡献 $\max\_k - k_0 + 1$ 个合法子方阵。

样例演示

样例 1：["red","red","red"], myval=2

矩阵为：

r e d
r e d
r e d

所有 $1 \times 1$ 方阵：只含一种字符，权值为 0（另外两种字符计数为 0），不满足 $\geq 2$ 。
所有 $2 \times 2$ 方阵：至多包含 $2$ 种字符各 $2$ 个，第 $3$ 种字符为 $0$ ，权值为 $0$ 。
$3 \times 3$ 方阵：r、e、d 各出现 $3$ 次，权值为 $3 \geq 2$ ，合法。

答案为 $1$ 。

样例 2：["red","edr","dre"], myval=1

矩阵为：

r e d
e d r
d r e

$1 \times 1$ 方阵权值为 $0$ ；但 $2 \times 2$ 和 $3 \times 3$ 方阵中许多都满足三种字符各至少出现 $1$ 次。答案为 $5$ 。

复杂度分析

时间复杂度： $O(n^2 \log n)$ 。枚举 $n^2$ 个左上角，每个做 $O(\log n)$ 的二分查找。
空间复杂度： $O(n^2)$ 。存储三个 $n \times n$ 的前缀和数组。

代码

C++
Java

class Solution {
public:
    int matrixCount(vector<string>& a, int myval) {
        int n = a.size();
        vector<vector<vector<int>>> pre(3, vector<vector<int>>(n + 1, vector<int>(n + 1, 0)));
        for (int i = 1; i <= n; i++) {
            for (int j = 1; j <= n; j++) {
                char ch = a[i - 1][j - 1];
                int idx = (ch == 'r') ? 0 : (ch == 'e') ? 1 : 2;
                for (int c = 0; c < 3; c++) {
                    pre[c][i][j] = pre[c][i - 1][j] + pre[c][i][j - 1]
                                 - pre[c][i - 1][j - 1] + (c == idx ? 1 : 0);
                }
            }
        }
        int ans = 0;
        for (int i = 1; i <= n; i++) {
            for (int j = 1; j <= n; j++) {
                int maxK = min(n - i + 1, n - j + 1);
                int lo = 1, hi = maxK, best = maxK + 1;
                while (lo <= hi) {
                    int mid = (lo + hi) / 2;
                    int r2 = i + mid - 1, c2 = j + mid - 1;
                    int minCnt = INT_MAX;
                    for (int c = 0; c < 3; c++) {
                        int cnt = pre[c][r2][c2] - pre[c][i - 1][c2]
                                - pre[c][r2][j - 1] + pre[c][i - 1][j - 1];
                        minCnt = min(minCnt, cnt);
                    }
                    if (minCnt >= myval) {
                        best = mid;
                        hi = mid - 1;
                    } else {
                        lo = mid + 1;
                    }
                }
                if (best <= maxK) {
                    ans += maxK - best + 1;
                }
            }
        }
        return ans;
    }
};

import java.util.*;

public class Solution {
    public int matrixCount(String[] a, int myval) {
        int n = a.length;
        int[][][] pre = new int[3][n + 1][n + 1];
        for (int i = 1; i <= n; i++) {
            for (int j = 1; j <= n; j++) {
                char ch = a[i - 1].charAt(j - 1);
                int idx = (ch == 'r') ? 0 : (ch == 'e') ? 1 : 2;
                for (int c = 0; c < 3; c++) {
                    pre[c][i][j] = pre[c][i - 1][j] + pre[c][i][j - 1]
                                 - pre[c][i - 1][j - 1] + (c == idx ? 1 : 0);
                }
            }
        }
        int ans = 0;
        for (int i = 1; i <= n; i++) {
            for (int j = 1; j <= n; j++) {
                int maxK = Math.min(n - i + 1, n - j + 1);
                int lo = 1, hi = maxK, best = maxK + 1;
                while (lo <= hi) {
                    int mid = (lo + hi) / 2;
                    int r2 = i + mid - 1, c2 = j + mid - 1;
                    int minCnt = Integer.MAX_VALUE;
                    for (int c = 0; c < 3; c++) {
                        int cnt = pre[c][r2][c2] - pre[c][i - 1][c2]
                                - pre[c][r2][j - 1] + pre[c][i - 1][j - 1];
                        minCnt = Math.min(minCnt, cnt);
                    }
                    if (minCnt >= myval) {
                        best = mid;
                        hi = mid - 1;
                    } else {
                        lo = mid + 1;
                    }
                }
                if (best <= maxK) {
                    ans += maxK - best + 1;
                }
            }
        }
        return ans;
    }
}