大小为 K 的子阵列中谷元素的最大计数数

问题描述

给定一个大小为 $n \times m$ 的矩阵，以及一个正整数 $k$，请你求出该矩阵中所有大小为 $k \times k$ 的子矩阵中，谷元素的最大计数数。谷元素是指在矩阵中行、列方向均没有比它更小的元素。

例如，对于下面的矩阵：

1 2 2 4
4 5 6 7
7 8 9 10
11 12 13 14

其中大小为 $2 \times 2$ 的子矩阵中，谷元素的最大计数数是 2，分别是：

1 2
4 5

和

5 6
8 9

解决方案

暴力枚举所有的子矩阵肯定是不行的，时间复杂度是 $O(n^2m^2)$，显然会 TLE。那么有没有比较高效的解法呢？

先考虑如何判断一个元素是否是谷元素。对于矩阵中的一个点 $(i, j)$，如果其在行方向和列方向上都是最小的，那么它就是一个谷元素。

接下来考虑如何高效找到每个子矩阵中的谷元素的数量。对于一个子矩阵，如果其中有谷元素，那么一定存在一行和一列，这行和这列上的所有元素都在子矩阵中。

于是可以先对每一行求出该行的最小值及其下标，然后再对每一列求出该列的最小值及其下标。

接着对于每个子矩阵，遍历其所有的行和列，统计其中的谷元素即可。

时间复杂度为 $O(nm(m+k))$。

代码实现

def count_valley_elements(matrix, k):
    n, m = len(matrix), len(matrix[0])
    row_min = [[0] * m for _ in range(n)]
    col_min = [[0] * m for _ in range(n)]
    for i in range(n):
        row_min_val, row_min_idx = 1000000000, -1
        for j in range(m):
            if matrix[i][j] < row_min_val:
                row_min_val, row_min_idx = matrix[i][j], j
            row_min[i][j] = row_min_idx
    for j in range(m):
        col_min_val, col_min_idx = 1000000000, -1
        for i in range(n):
            if matrix[i][j] < col_min_val:
                col_min_val, col_min_idx = matrix[i][j], i
            col_min[i][j] = col_min_idx
    ans = 0
    for i in range(n - k + 1):
        for j in range(m - k + 1):
            cnt = 0
            for p in range(k):
                for q in range(k):
                    if row_min[i + p][j + q] != j + q or col_min[i + p][j + q] != i + p:
                        break
                else:
                    cnt += 1
            ans = max(ans, cnt)
    return ans

代码中用到的两个二维数组 row_min 和 col_min 分别记录了矩阵中每一行和每一列的最小值及其下标。

最后对于每个子矩阵，使用双重循环遍历其所有的行和列，统计其中的谷元素即可。