Question

I’m currently developing stochastic optimization algorithms and have encountered the following issue (which I imagine appears also in other places): It could be called totally unstable partial sort:

Given a container of size n and a comparator, such that entries may be equally valued.
Return the best k entries, but if values are equal, it should be (nearly) equally probable to receive any of them.

(output order is irrelevant to me, i.e. equal values completely among the best k need not be shuffled. To even have all equal values shuffled is however a related, interesting question and would suffice!)

A very (!) inefficient way would be to use shuffle_randomly and then partial_sort, but one actually only needs to shuffle the block of equally valued entries “at the selection border” (resp. all blocks of equally valued entries, both is much faster). Maybe that Observation is where to start…

I would very much prefer, if someone could provide a solution with STL algorithms (or at least to a large portion), both because they’re usually very fast, well encapsulated and OMP-parallelized.

Thanx in advance for any ideas!

Answer 1

If you really mean that output order is irrelevant, then you want std::nth_element, rather than std::partial_sort, since it is generally somewhat faster. Note that std::nth_element puts the n^th element in the right position, so you can do the following, which is 100% standard algorithm invocations (warning: not tested very well; fencepost error possibilities abound):

template<typename RandomIterator, typename Compare>
void best_n(RandomIterator first,
            RandomIterator nth,
            RandomIterator limit,
            Compare cmp) {
  using ref = typename std::iterator_traits<RandomIterator>::reference;
  std::nth_element(first, nth, limit, cmp);
  auto p = std::partition(first, nth, [&](ref a){return cmp(a, *nth);});
  auto q = std::partition(nth + 1, limit, [&](ref a){return !cmp(*nth, a);});
  std::random_shuffle(p, q);  // See note
}

The function takes three iterators, like nth_element, where nth is an iterator to the n^th element, which means that it is begin() + (n - 1)).

Edit: Note that this is different from most STL algorithms, in that it is effectively an inclusive range. In particular, it is UB if nth == limit, since it is required that *nth be valid. Furthermore, there is no way to request the best 0 elements, just as there is no way to ask for the 0^th element with std::nth_element. You might prefer it with a different interface; do feel free to do so.

Or you might call it like this, after requiring that 0 < k <= n:

best_n(container.begin(), container.begin()+(k-1), container.end(), cmp);

It first uses nth_element to put the “best” k elements in positions 0..k-1, guaranteeing that the k^th element (or one of them, anyway) is at position k-1. It then repartitions the elements preceding position k-1 so that the equal elements are at the end, and the elements following position k-1 so that the equal elements are at the beginning. Finally, it shuffles the equal elements.

nth_element is O(n); the two partition operations sum up to O(n); and random_shuffle is O(r) where r is the number of equal elements shuffled. I think that all sums up to O(n) so it’s optimally scalable, but it may or may not be the fastest solution.

---

Note: You should use std::shuffle instead of std::random_shuffle, passing a uniform random number generator through to best_n. But I was too lazy to write all the boilerplate to do that and test it. Sorry.