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

题意

给定非递减排序的整数数组 nums 和整数 target，返回 target 在数组中的起止下标 [first, last]；不存在则返回 [-1, -1]。

要求时间复杂度 $O(\log n)$ 。

思路

线性扫描可以找到答案，但需要 $O(n)$ 时间。利用数组有序性，可以用两次二分将复杂度降到 $O(\log n)$ ：

先看一个可以直接验证想法的朴素解：

cpp

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// brute.cpp：小数据暴力解，线性扫描第一个和最后一个等于 target 的位置。
#include <bits/stdc++.h>
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, t;
    cin >> n >> t;
    vector<int> a(n);
    for (int &x : a)
        cin >> x;

    int first = -1, last = -1;
    // 从左到右扫一遍，记录 target 第一次和最后一次出现的位置。
    for (int i = 0; i < n; i++) {
        if (a[i] == t) {
            if (first == -1)
                first = i; // 第一次命中
            last = i;      // 持续更新到最后一次
        }
    }
    cout << first << ' ' << last << '\n';
    return 0;
}

朴素解线性扫描整个数组，时间 $O(n)$ ，不满足进阶要求。

优化的关键是：lower_bound 返回第一个 $\geqslant \text{target}$ 的位置，upper_bound 返回第一个 $> \text{target}$ 的位置。若 lower_bound 所指元素恰好等于 target，则 lower_bound 就是起始下标，upper_bound - 1 就是终止下标；若不等于，说明 target 不存在，返回 [-1, -1]。

判断条件 l > r（即 lower_bound > upper_bound - 1）统一覆盖了空数组、target 不存在和 target 越界三种情况。

Python 实现用同一个 bs 函数，参数 left=True 时行为同 lower_bound（ $\geqslant$ ），left=False 时行为同 upper_bound（ $>$ ）。谓词 nums[m] > target or (left and nums[m] == target) 在左界模式下把等于也视为"往左缩"，右界模式下只有严格大于才缩左。

代码

cpp

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#include <bits/stdc++.h>
using namespace std;

class Solution {
public:
    vector<int> searchRange(vector<int> &nums, int target) {
        int l = lower_bound(nums.begin(), nums.end(), target) - nums.begin();
        int r = upper_bound(nums.begin(), nums.end(), target) - nums.begin() - 1;
        if (l > r)
            return {-1, -1};
        return {l, r};
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n, t;
    cin >> n >> t;
    vector<int> a(n);
    for (int &x : a)
        cin >> x;
    auto v = Solution().searchRange(a, t);
    cout << v[0] << ' ' << v[1] << '\n';
    return 0;
}

python

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
30#!/usr/bin/env python3
from typing import List


class Solution:
    def searchRange(self, nums: List[int], target: int) -> List[int]:
        def bs(lo, hi, left):
            while lo <= hi:
                m = (lo + hi) // 2
                if nums[m] > target or (left and nums[m] == target):
                    hi = m - 1
                else:
                    lo = m + 1
            return lo

        n = len(nums)
        l = bs(0, n - 1, True)
        if l == n or nums[l] != target:
            return [-1, -1]
        return [l, bs(0, n - 1, False) - 1]


def main():
    n, t = map(int, input().split())
    a = list(map(int, input().split())) if n else []
    print(*Solution().searchRange(a, t))


if __name__ == "__main__":
    main()

复杂度

时间复杂度： $O(\log n)$ ，两次二分各扫描一次，每次比较 $O(1)$ 。
空间复杂度： $O(1)$ ，只用到常数个变量。

总结

在有序数组上查找边界，核心是把"找第一个/最后一个等于"转化为标准二分谓词：左界用 $\geqslant$ （lower_bound），右界用 $>$ （upper_bound），两者之差即为出现区间长度。l > r 是判断不存在最简洁的条件。

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

目录

题意

思路

代码

复杂度

总结