1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55
| #include <iostream> #include <vector> #include <algorithm> #include <cstdlib>
template <typename RandomIt> RandomIt partition(RandomIt first, RandomIt last, RandomIt pivot) { std::iter_swap(pivot, last - 1); RandomIt storeIndex = first;
for (RandomIt it = first; it < last - 1; ++it) { if (*it < *(last - 1)) { std::iter_swap(it, storeIndex); ++storeIndex; } } std::iter_swap(storeIndex, last - 1); return storeIndex; }
template <typename RandomIt> void quickselect(RandomIt first, RandomIt last, size_t n) { if (first < last) { RandomIt pivot = first + std::rand() % (last - first); pivot = partition(first, last, pivot);
if (pivot - first == n) { return; } else if (pivot - first > n) { quickselect(first, pivot, n); } else { quickselect(pivot + 1, last, n - (pivot - first + 1)); } } }
template <typename RandomIt> void my_nth_element(RandomIt first, RandomIt nth, RandomIt last) { size_t n = nth - first; quickselect(first, last, n); }
int main() { std::vector<int> vec = {3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5}; size_t n = 5;
my_nth_element(vec.begin(), vec.begin() + n, vec.end());
std::cout << "The " << n << "th element is: " << vec[n] << std::endl;
return 0; }
|