题目

在排序数组中查找元素的第一个和最后一个位置

思路

数组是排序数组，要求$target$值的最小和最大下标，直接套用二分法模板。

以求$target$的最小下标为例:

  1. 搜索区间$[left, rigt]$的闭区间，初始值为$[0, len(nums)-1]$
  2. $nums[mid] = target$，则区间左移，$right=mid-1$
  3. $nums[mid] \lt target$，则区间左移，$right=mid-1$
  4. $nums[mid] \gt target$，则区间右移，$left=mid+1$

其实只要仔细分析每次搜索区间的缩小条件，问题就会变得很简单。

代码

完全按照思路分析中的步骤，用代码实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class Solution:
    def searchRange(self, nums: List[int], target: int) -> List[int]:
        n = len(nums)
        if n == 0:
            return [-1, -1]
        left = 0
        right = n - 1

        # 求左边界
        while left <= right:
            mid = left + (right - left) // 2
            if nums[mid] >= target:
                right = mid - 1
            else:
                left = mid + 1
        lr = left if left < n and nums[left] == target else -1

        # 求右边界
        left = 0
        right = n - 1
        while left <= right:
            mid = left + (right - left) // 2
            if nums[mid] > target:
                right = mid - 1
            else:
                left = mid + 1
        rr = right if right < n and nums[right] == target else -1
        return [lr, rr] if nums[lr] == nums[rr] else [-1, -1]

优化代码

可以看到上述代码有很多重复的部分，咱们可以稍微优化一下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

class Solution:
    def searchRange(self, nums: List[int], target: int) -> List[int]:
        n = len(nums)
        if n == 0:
            return [-1, -1]

        def _search(is_lower=False):
            left, right = 0, n - 1
            while left <= right:
                mid = left + (right - left) // 2
                if nums[mid] == target:
                    if is_lower:
                        right = mid - 1
                    else:
                        left = mid + 1
                elif nums[mid] > target:
                    right = mid - 1
                else:
                    left = mid + 1
            res = left if is_lower else right
            return res if res < n and nums[res] == target else -1

        return [_search(is_lower=True), _search()]