B - Trimmed Mean Editorial by MMNMM

より高速な解法

quick select と呼ばれるアルゴリズムを用いて \(O(N)\) 時間で解く解法について解説します。

quick select について


quick select は、長さ \(N\) の配列 \((A[i]) _ {1\leq i\leq N}\) と整数 \(k\ (1\leq k\leq N)\) を入力として、\(A[i]\) を並べ替えた列 \(A^\prime[i]\) であって次の条件を満たすものを \(O(N)\) 時間で得るアルゴリズムです。

  • \(i\leq k\iff A^\prime[i]\leq A^\prime[k]\quad(1\leq i\leq N)\)
  • \(i\geq k\iff A^\prime[i]\geq A^\prime[k]\quad(1\leq i\leq N)\)

単に

  • \(A^\prime[k]=A[i]\ (1\leq i\leq N)\) の中で \(k\) 番目に大きいもの

を得るアルゴリズムだと考えることが多いですが、一般的な実装では結果的にこのような配列が得られます。 (また、単に \(A^\prime[k]\) の値だけが得られる場合でも、クイックソートのピボットで列を分割する要領で上のような配列 \(A^\prime\) が追加の \(O(N)\) 時間をかけることで得られます。)

quick select は、クイックソートを行う際に \(k\) 番目の要素が含まれない側のソートを行わないことで実現できます。 時間計算量は、ピボットの選択に乱択を用いると期待 \(O(N)\) 時間、median of medians アルゴリズムなどを用いると最悪 \(O(N)\) 時間となります。

C++ には、これを行う(厳密には、少し異なるアルゴリズムの introselect を行う) std::nth_element という関数があり、平均 \(O(N)\) 時間であることが保証されています。事前に \(O(N)\) 時間で配列をシャッフルしておくことで期待 \(O(N)\) 時間とすることもできます。

この問題を解く


長さ \(5N\) の列 \((X[i]) _ {1\leq i\leq 5N}\) と \(N\) に対して quick select を行うと、大きいほうから \(4N\) 人ぶん(順不同)の列 \((Y[i]) _ {1\leq i\leq 4N}\) が得られます。 この列と \(3N\) に対して quick select を行うと、この中で大きくないほうから \(3N\) 人ぶん(順不同)の列 \((Z[i]) _ {1\leq i\leq 3N}\) が得られます。

出力すべき値は \(\displaystyle\dfrac1{3N}\sum _ {i=1} ^ {3N}Z[i]\) です。 時間計算量は \(O(N)\) 時間です。

実装例は以下のようになります。 なお、誤差については浮動小数点で正確に表せる値どうしの除算なので、誤差は \(0.5\operatorname{ulp}\) 以下となり十分小さいです(例えば計算に float 型を用いた場合、相対誤差は \(1.2\times10^{-7}\) 以下と示せます)。

#include <iostream>
#include <vector>
#include <algorithm>
#include <numeric>

int main() {
    using namespace std;

    unsigned N;
    cin >> N;
    vector<unsigned> A(5 * N);
    for(auto&& a : A)cin >> a;

    nth_element(begin(A), begin(A) + N, end(A)); // 小さいほうから N 個を除く
    nth_element(begin(A) + N, end(A) - N, end(A)); // 大きいほうから N 個を除く

    cout << reduce(begin(A) + N, end(A) - N) / static_cast<long double>(3 * N) << endl;

    return 0;
}

posted:
last update: