Keep smallest n elements in set, discard others

Question

I have a C++ STL set.

set<unsigned> baseSet;

I insert a bunch of numbers into it.

for(auto incoming : it->second.in) {
    baseSet.insert(incoming);
}

Now I want only the smallest n elements in the set. (I can switch the set's comparator to std::greater if that helps.) Sets don't seem to have a resize function like vectors. How can I shorten it? In-place would be great, and I definately want to take advantage of the fact that sets are sorted instead of doing all the sorting again.

The node based containers such as std::set are for specialized uses and you typically do not need and should not use them. You'd probably do better to use a vector and nth_element. — bames53
Unlike vectors, sets are made specifically for the case where you don't want repeated elements, which is exactly what I need for my application. — Qaz
Even so, a vector is probably better. Just use std::unique. See Why you shouldn't use set (and what you should use instead) and Why you should avoid Linked Lists. (The second link refers to linked lists, but the information can also be applied to std::set.) — bames53
The point @bames53 is making is that constructing the set has complexity O(N lg K), where N is the total number of insertions and K the number of distinct elements. Throwing everything in a vector which you then sort, unique and resize to n has complexity O(N lg N) but much more favorable constant factors. If K is not significantly smaller than N, or N is not "large", using a vector will almost certainly be faster. — Casey

Mohit Jain Mohit Jain · Accepted Answer · 2014-11-18T04:46:11

Option 1

if(baseSet.size() > n) {  // Use appropriate case if necessary
  auto it = base.begin();  // std::set<unsigned>::iterator
  std::advance (it, n);
  baseSet.erase(it, baseSet.end());
}

Find n^th smallest element and erase range from that element to end.

Option 2

c++11 onwards, you can also do this as following (Thanks to bames53 for the comment)

baseSet.erase(next(begin(baseSet), n), end(baseSet));

Option 3

Another option (less recommended) is to keep deleting the elements once set reaches the size n.

for(auto incoming : it->second.in) {
    baseSet.insert(incoming);  // Insert conditionally if you bother about efficiency
    if(baseSet.size() > n) baseSet.erase(baseSet.rbegin());  // rbegin not to be confused with begin
}

Keep smallest n elements in set, discard others

1 Answers