k_th_finder.cc

const size_t K_TH_FIND_LENGTH_THRESHOLD = 5;

template <typename T>
bool k_th_find_default_comparator(T l, T r) {
    return l < r;
}

template <typename T>
void k_th_find_partition(std::vector<T> *_v, T e,                                       //
                         size_t *_idx,                                                  //
                         bool (*comparator)(T, T) = k_th_find_default_comparator<T>) {  //
    size_t pos = find(_v->begin(), _v->end(), e) - _v->begin();
    size_t last_pos = _v->size() - 1;

    std::iter_swap(_v->begin() + pos, _v->begin() + last_pos);
    size_t i = 0;
    for (size_t j = 0; j < last_pos; j++) {
        if (comparator(*(_v->begin() + j), e)) {
            std::iter_swap(_v->begin() + i, _v->begin() + j);
            i++;
        }
    }
    std::iter_swap(_v->begin() + i, _v->begin() + last_pos);

    *_idx = i;
}

template <typename T>
T k_th_find(std::vector<T> v,                                              //
            size_t k,                                                      //
            bool (*comparator)(T, T) = k_th_find_default_comparator<T>) {  //
    size_t len = v.size();
    if (len > K_TH_FIND_LENGTH_THRESHOLD) {
        std::vector<T> vn;
        for (size_t i = 0; i < len; i += K_TH_FIND_LENGTH_THRESHOLD) {
            size_t lmt = i + K_TH_FIND_LENGTH_THRESHOLD < len ? i + K_TH_FIND_LENGTH_THRESHOLD : len;
            std::sort(&v[i], &v[lmt], comparator);
            vn.push_back(v[(lmt + i - 1) / 2]);
        }

        T e = k_th_find(vn, (vn.size() - 1) / 2, comparator);
        size_t idx = 0;
        k_th_find_partition(&v, e, &idx, comparator);
        if (idx == k) {
            return e;
        } else if (idx > k) {
            std::vector<T> v_sub(v.begin(), v.begin() + idx);
            return k_th_find(v_sub, k, comparator);
        } else {
            std::vector<T> v_sub(v.begin() + idx + 1, v.end());
            return k_th_find(v_sub, k - idx - 1, comparator);
        }
        // return e;
    } else {  // length shall > 0..
        std::sort(v.begin(), v.end(), comparator);

        return v[k];
    }
}