# LeetCode Medium 658. Find K Closest Elements

tags: leetcode

問題

• Explore
• Problems

アイデア

整列済みの配列とkとxが与えられた時に，xに最も近いk個の要素を見つけるという問題．返り値の配列は昇順で出来るだけ小さい要素から入れます．

整列済みなので二分探索が使えそうです．

解法

x以上の最小の要素を求めます．この要素よりも左k個分左のインデックス(配列外ならば0番目)，この要素よりも右にk-1個分右のインデックスを取ります．これが探索する左端と右端です．

そして右端と左端の要素のうちxとの差が小さい方を答えとします．(つまり，大きい方の端を削ります)

計算量

配列の大きさを$N$として

• 時間計算量
  • 最小の処理で$O(\log N)$必要です．
  • そして，返す配列の要素数を減らす処理で$O(K)$必要です
  • 合計で$O(\log N + K)$
• 空間計算量
  • これは返す配列をコピーするので$O(K)$

C++

class Solution {
public:
  vector<int> findClosestElements(vector<int>& arr, int k, int x) {
    int n = arr.size();
    int idx = lower_bound(arr.begin(), arr.end(), x) - arr.begin();
    int left = max(0, idx - k);
    int right = min(n - 1, idx + k - 1);
    while (right - left + 1 > k) {
      if (x - arr[left] > arr[right] - x) {
        left++;
      } else {
        right--;
      }
    }
    return vector<int>(arr.begin() + left, arr.begin() + right + 1);
  }
};