C++で配列内のk個の最強値を求めるアルゴリズム

問題の概要

数値の配列 arr と整数 k が与えられているとします。ここで、値 arr[i] が値 arr[j] より「強い(strong)」とは、次の条件を満たすことを指します。

|arr[i] − m| > |arr[j] − m|(m は配列の中央値)

また、|arr[i] − m| と |arr[j] − m| が等しい場合は、arr[i] > arr[j] であるときに arr[i] の方が強いとみなされます。このルールに基づき、配列の中から最も強い k 個の値をリストとして返すのが本問題の目的です。

具体例

入力が arr = [1,2,3,4,5]、k = 2 の場合を考えてみましょう。このとき中央値は 3 であり、強さの順に要素を並べ替えると [5, 1, 4, 2, 3] となります。したがって、最も強い 2 つの要素は [5, 1](または [1, 5])です。

|5 − 3| と |1 − 3| はどちらも 2 で等しいのですが、5 > 1 であるため、5 の方が 1 より強いと判定されます。

解法のアプローチ

この問題は、ソート + 双方向ポインタ(ツーポインタ)のテクニックを使うことで効率的に解けます。手順は以下の通りです。

• 配列 arr をソートする

• n := arr のサイズとする

• m := arr[(n − 1) / 2](ソート後の中央値)

• 結果格納用の配列 ret を定義する

• i := 0(先頭ポインタ)、j := n − 1(末尾ポインタ)とする

• k が 0 になるまで以下を繰り返す:

  • x1 := |arr[j] − m|(末尾側の中央値からの距離)

  • x2 := |arr[i] − m|(先頭側の中央値からの距離)

  • x1 ≥ x2 の場合:ret の末尾に arr[j] を追加し、j を 1 減らす

  • それ以外の場合:ret の末尾に arr[i] を追加し、i を 1 増やす

• ret を返す

ソート済みの配列では、中央値から最も遠い候補は必ず先頭か末尾に存在します。そのため、両端を比較しながら強い方を取り出していくこの方法で、毎回正しく最強の値を選択できます。計算量はソートに O(n log n)、抽出処理に O(k) となり、非常に効率的です。

C++での実装例

以下のコードで実際の動作を確認できます。

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<int> v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
        cout << v[i] << ", ";
    }
    cout << "]"<<endl;
}
class Solution {
public:
    int calc(int x, int m){
        return abs(x - m);
    }
    vector<int> getStrongest(vector<int>& arr, int k) {
        sort(arr.begin(), arr.end());
        int n = arr.size();
        int m = arr[(n - 1) / 2];
        vector<pair<int, int> > v;
        int i = 0;
        int j = n - 1;
        vector<int> ret;
        while (k--) {
            int x1 = calc(arr[j], m);
            int x2 = calc(arr[i], m);
            if (x1 >= x2) {
                ret.push_back(arr[j]);
                j--;
            }
            else {
                ret.push_back(arr[i]);
                i++;
            }
        }
        return ret;
    }
};
main(){
    Solution ob;
    vector<int> v = {1,2,3,4,5};
    print_vector(ob.getStrongest(v,2));
}

入力

{1,2,3,4,5},2

出力

[5, 1]

まとめ

本記事では、中央値からの距離と値の大小を基準に「強さ」を判定し、配列から最強の k 個の値を抽出する方法を紹介しました。ポイントは、配列をソートした上で両端から比較を進める双方向ポインタの活用です。同点の場合は大きい値を優先するという条件(x1 ≥ x2 の分岐)を正しく実装することが、正解への鍵となります。