【C++】二分探索でxに最も近いK個の要素を効率的に求める方法

問題概要

ソート済みの整数配列と、2つの整数 k および x が与えられたとします。このとき、配列の中から x に最も近い k 個の要素を求めます。結果は昇順で返す必要があり、距離が等しい要素が存在する場合は、より小さい値の要素を優先します。

例えば、入力が [1,2,3,4,5]、k = 4、x = 3 の場合、出力は [1,2,3,4] となります。

アルゴリズムの考え方（二分探索）

この問題は二分探索を活用することで、O(log(n − k) + k) の計算量で効率的に解くことができます。ポイントは、「答えとなる k 個の連続した部分配列の開始位置」を二分探索で絞り込むことです。

候補となるウィンドウの左端 mid に対して、x − arr[mid]（左端との距離）が arr[mid + k] − x（右端との距離）より大きい場合は、ウィンドウを右へずらすべきなので low = mid + 1 とします。そうでなければ、答えは mid またはそれより左側にあるため high = mid とします。この条件式により、距離が同点の場合には自動的に小さい要素側が優先されます。

手順

• 結果を格納する配列 ans を用意する
• low = 0、high = 配列のサイズ − k と初期化する
• low < high の間、以下を繰り返す：
  • mid = low + (high − low) / 2 を計算する
  • x − arr[mid] > arr[mid + k] − x ならば low = mid + 1、そうでなければ high = mid
• i を low から low + k − 1 まで動かしながら、arr[i] を ans に追加する
• ans を返す

C++での実装例

以下の実装を見ると、より理解が深まるでしょう。

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
        cout << v[i] << ", ";
    }
cout << "]"<<endl;
}
class Solution {
public:
    vector<int> findClosestElements(vector<int>& arr, int k, int x) {
        vector <int> ans;
        int low = 0;
        int high = arr.size() - k;
        while(low < high){
            int mid = low + (high - low)/2;
            if(x - arr[mid] > arr[mid + k] - x){
                low = mid + 1;
            }
            else high = mid;
        }
        for(int i = low ; i < low + k ; i++)ans.push_back(arr[i]);
        return ans;
    }
};
main(){
    Solution ob;
    vector<int> v = {1,2,3,4,5};
    print_vector(ob.findClosestElements(v, 4, 3));
}

入力

[1,2,3,4,5]
4
3

出力

[1,2,3,4]

計算量の評価

時間計算量は、二分探索部分が O(log(n − k))、結果配列の構築が O(k) であるため、全体で O(log(n − k) + k) となります。単純に全要素の絶対差を計算してソートする方法（O(n log n)）と比べても、非常に効率的なアプローチだと言えます。補助的に使用するメモリは出力を除き O(1) です。