【算法】在二维不单调的矩阵上二分查找——力扣1901. 寻找峰值 II

1901. 寻找峰值 II

---

【算法】在二维不单调的矩阵上二分查找——力扣1901. 寻找峰值 II

问题描述

给定一个从0开始编号的m x n矩阵mat，其中任意两个相邻格子的值都不相同。峰值是指那些严格大于其相邻格子(上、下、左、右)的元素。需要找出任意一个峰值mat[i][j]并返回其位置[i, j]。

示例

示例 1:

输入: mat = [[1, 4], [3, 2]]
输出: [0, 1]
解释: 3 和 4 都是峰值，所以[1, 0]和[0, 1]都是可接受的答案。

示例 2:

输入: mat = [[10, 20, 15], [21, 30, 14], [7, 16, 32]]
输出: [1, 1]
解释: 30 和 32 都是峰值，所以[1, 1]和[2, 2]都是可接受的答案。

解决思路

步骤一：列转行

首先，将矩阵的列转换为行，表示为col[mid]代表原数组第mid列。我们要搜索出col[mid]，也就是第mid列的最大值的索引max_idx，以及这一列的最大值max_val。

步骤二：回到一维数组上的寻找峰值的思路

现在，只需要确保col[mid][max_idx]大于col[mid - 1][max_idx]和col[mid + 1][max_idx]这两个元素。显然，问题的形态与162.寻找峰值非常相似。现在问题转换成了类似在一维数组上的寻找峰值问题的形式。

步骤三：二分搜索

此时，可以使用二分搜索的思路。我们定义le_val = col[mid - 1][max_idx]，如果le_val < max_val，说明第mid列有可能包含峰值，也可能是在峰值的列的左侧，此时可以更新left = mid；如果le_val >= max_val，说明mid必定不是峰值，而左边的列有可能是峰值，此时更新right = mid - 1。

代码实现

class Solution:
    def findPeakGrid(self, mat: List[List[int]]) -> List[int]:
        m, n = len(mat), len(mat[0])
        cols = list(zip(*mat))

        def get_val(x, y):
            if not (0 <= x < n and 0 <= y < m):
                return -1
            return cols[x][y]

        left, right = 0, n - 1  # n列
        while left < right:
            mid = (left + right + 1) // 2
            max_idx, max_val = max(enumerate(cols[mid]), key=lambda v: v[1])
            le_val = get_val(mid - 1, max_idx)  # 与左边的行作比较
            if le_val < max_val:
                left = mid
            else:
                right = mid - 1

        ans_col = left
        ans_row = max(enumerate(cols[ans_col]), key=lambda v: v[1])[0]
        return [ans_row, ans_col]

二分示意图

二分初始的状态

如图所示，红色的20是mid所指向的列的最大值，获取这个最大值的索引以及值的相关的代码是：

max_idx, max_val = max(enumerate(cols[mid]), key=lambda v: v[1]) 

left		mid		right
1	2	3	12	13
4	5	6	14	15
7	8	9	16	17
18	19	20	21	22

图1. 二分初始的状态

二分更新说明

因为20比左边的19大，满足le_val < max_val，说明mid所在的这一列或其右侧的列包含峰值，所以，更新left = mid后的情况如下图所示

二分更新后的状态

		left	mid	right
1	2	3	12	13
4	5	6	14	15
7	8	9	16	17
18	19	20	21	22

图2. 二分更新后的状态

性能分析

因为一共有n列，m行，我们是二分查找峰值所在的列，每一次都要遍历一整列上，以找到最大值，一整列上有m个元素，所以，该算法的时间复杂度为O(m log(n))，满足题目要求。