k番目要素

数列が与えられ，k番目の要素を求めるという問題がある．

解き方としてクイックソートに類似したものがあり，計算時間の期待値は O(n) である．

参考：Spaghetti Source - k 番目の要素の選択

この問題をとく乱択でないアルゴリズムをしった（上記サイトにも解説がある）ので，練習のため実装した． クイックソートライクのアルゴリズムと似ているが，ピボット選択が悪くならないように工夫をしている．

以下コード．ここではkを0はじまりとしている．

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

template <class T>
T calc_k_th_smallest_elem(const std::vector<T>& v, size_t k){
  // k in [0, #v)
  assert(k<v.size());

  // Base case, #v <= 10
  if(v.size() <= 10){
    auto v_copied(v);
    std::sort(v_copied.begin(), v_copied.end());
    return v_copied[k];
  }

  // #v > 10
  std::vector<T> mids;
  mids.reserve(v.size()/5);
  for(size_t i=0; i<v.size()/5; i+=5){
    std::vector<T> mid_candidates;
    for(size_t j=0; j<5; ++j){mid_candidates.push_back(v[i+j]);}
    sort(mid_candidates.begin(), mid_candidates.end());
    mids.push_back(mid_candidates[2]);
  }

  T m = calc_k_th_smallest_elem(mids, mids.size()/2);
  std::vector<T> lt_m;
  std::copy_if(v.begin(), v.end(), std::back_inserter(lt_m), [&m](const T& x)->bool{return x<m;});
  std::vector<T> eq_m;
  std::copy_if(v.begin(), v.end(), std::back_inserter(eq_m), [&m](const T& x)->bool{return x==m;});
  std::vector<T> gt_m;
  std::copy_if(v.begin(), v.end(), std::back_inserter(gt_m), [&m](const T& x)->bool{return x>m;});

  assert(gt_m.size() + eq_m.size() + lt_m.size() == v.size());

  if(lt_m.size() > k){
    return calc_k_th_smallest_elem(lt_m, k);
  }
  else if (lt_m.size() + eq_m.size() > k){
    return m;
  }
  else{
    return calc_k_th_smallest_elem(gt_m, k - lt_m.size() - eq_m.size());
  }
}

clang -O2 でコンパイルすると， n = 1億, T = int で6+秒の実行時間でした． Spaghetti Source - k 番目の要素の選択 にもあるように，クイックソート版より実装は複雑である．が，C++11の恩恵もあって，それなりに読める長さになっていると思う．特にラムダ式とcopy_ifが貢献している…気がします．

このコードは「データ構造とアルゴリズム（杉原）」の本アルゴリズムの解説をもとに書き起こしたものです．

データ構造とアルゴリズム