LC658.找到K个最接近的元素

原题链接
image.png

思路

  • 二分查找 + 双指针
  • 先利用二分查找，找到>= x的第一个值（>= x区间的左端点），记为pos。但是pos位置的值，未必是离x最近的值。可以在排除特殊情况pos == 0 || pos == len - 1之后，比较arr[pos - 1]和arr[pos]离x哪个更近，更近的那个更新为pos位置。如果一样近，就选择左边的那个更新成pos位置。

  • 左指针lp，右指针rp，起始时刻都指向离x最近的位置pos。然后左右移动两个指针，使得区间长度恰好为k，然后返回数组切片arr[lp:rp]即可。这一点看似简单，其实细节极其繁琐，不好弄。需要关注的细节如下。

    • 特殊情况：如果左指针到达边界lp == 0，那么直接动右指针即可；如果右指针到达边界rp == len - 1，那么直接移动左指针即可。
    • 一般情况下的移动规则：左右两根指针的下一个位置arr[lp - 1]和arr[rp + 1]，哪一个离x近，哪一个就往边缘动一格。

    • 最后的终止条件，lp和rp位置都是需要取进去的，所以长度rp - lp + 1 == k，就说明满足长度条件，从而退出循环。

      代码

      1. class Solution {
      2. public:
      4. int findPos(vector<int>& arr, int x) {
      5.    int left = 0, right = arr.size() - 1;
      6.    while (left < right) {
      7.        int mid = (left + right) / 2;
      8.        if (arr[mid] >= x) {
      9.            right = mid;
      10.        } else {
      11.            left = mid + 1;
      12.        }
      13.    }
      15.    if (arr[left] >= x) {
      16.        return left;
      17.    } else {
      18.        return -1;
      19.    }
      20. }
      22. vector<int> findClosestElements(vector<int>& arr, int k, int x) {
      24.    // 找到大于等于 x 的第一个数的下标
      25.    int n = arr.size();
      26.    int pos = findPos(arr, x);
      27.    // cout << pos << endl;
      29.    if (pos == -1) {
      30.        return {arr.end() - k, arr.end()};
      31.    } else if (pos == 0) {
      32.        return {arr.begin(), arr.begin() + k};
      33.    }
      35.    // 确保 pos 是离得最近的点 
      36.    // 如果距离一致，优先选左边
      37.    pos = (abs(arr[pos] - x) < abs(arr[pos - 1] - x) ? pos : pos - 1);
      39.    int lp = pos, rp = pos;
      41.    while (rp - lp + 1 < k) {
      42.        if (lp == 0) {
      43.            rp += 1;
      44.            continue;
      45.        } else if (rp == n - 1) {
      46.            lp -= 1;
      47.            continue;
      48.        }
      50.        if (abs(x - arr[lp - 1]) <= abs(arr[rp + 1] - x)) {
      51.            lp -= 1;
      52.        } else {
      53.            rp += 1;
      54.        }
      55.    }
      57.    // [lp, rp]包头包尾
      58.    return {arr.begin() + lp, arr.begin() + rp + 1};
      59. }
      60. };