C++: Get K smallest elements+indices from vector with ties - c++

Question

The task is to extract k smallest elements and their indices from double array, possibly including more elements that are tied to the k-th smallest one. E.g.:
input: {3.3,1.1,6.5,4.2,1.1,3.3}
output (k=3): {1,1.1} {4,1.1} {0,3.3} {5,3.3}
[This seems like a pretty common task, but I couldn't find a similar thread on SO - which handles ties. Hopefully, I didn't miss any and didn't duplicate the question.]
I came up with the following solution, which works and seems to be fairly efficient complexity-wise. E.g. for random 1MLN doubles and k=10 it takes ~40ms with MSVC 2013. I wonder if there's a better/cleaner/more efficient(for large data and/or large k) way to perform this task (validations for k value and similar things are our of scope here). Avoid allocating the queue with all elements? Make use of std::partial_sum or std::nth_element?
typedef std::pair<double, int> idx_pair;
typedef std::priority_queue<idx_pair, std::vector<idx_pair>, std::greater<idx_pair>> idx_queue;
std::vector<idx_pair> getKSmallest(std::vector<double> const& data, int k)
{
idx_queue q;
{
std::vector<idx_pair> idxPairs(data.size());
for (auto i = 0; i < data.size(); i++)
idxPairs[i] = idx_pair(data[i], i);
q = idx_queue(std::begin(idxPairs), std::end(idxPairs));
};
std::vector<idx_pair> result;
auto topPop = [&q, &result]()
{
result.push_back(q.top());
q.pop();
};
for (auto i = 0; i < k; i++)
topPop();
auto const largest = result.back().first;
while (q.empty() == false)
{
if (q.top().first == largest)
topPop();
else
break;
}
return result;
}
Working example is here.

Answer 1

Here's an alternative solution, suggested by #piotrekg2 - using nth_element with average O(N) complexity:
bool equal(double value1, double value2)
{
return value1 == value2 || std::abs(value2 - value1) <= std::numeric_limits<double>::epsilon();
}
std::vector<idx_pair> getNSmallest(std::vector<double> const& data, int n)
{
std::vector<idx_pair> idxPairs(data.size());
for (auto i = 0; i < data.size(); i++)
idxPairs[i] = idx_pair(data[i], i);
std::nth_element(std::begin(idxPairs), std::begin(idxPairs) + n, std::end(idxPairs));
std::vector<idx_pair> result(std::begin(idxPairs), std::begin(idxPairs) + n);
auto const largest = result.back().first;
for (auto it = std::begin(idxPairs) + n; it != std::end(idxPairs); ++it)
if (equal(it->first, largest))
result.push_back(*it);
return result;
}
Indeed, the code looks a bit cleaner. However, I've run some tests and empirically this solution is slightly slower than the original one with std::priority_queue.
Note: The answer below by Petar offers a similar solution using std::nth_element, which in my experiments, performs slightly better than this one and also better than the solution using std::priority_queue - perhaps because of eliminating the operation on pairs and working with primitive doubles instead.

Answer 2

As pointed out by asker, I will suggest first copy the vector of double and use a nth_element to find out the kth element.
Then do a linear scan and get the elements that are smaller than or equal to the kth element. The Time complexity should be linear.
However, it should be careful when comparing double.
vector<idx_pair> getKSmallest(vector<double> const& data, int k){
vector<double> data_copy = data;
nth_element(data_copy.begin(), data_copy.begin() + k, data_copy.end());
vector<idx_pair> result;
double kth_element = data_copy[k - 1];
for (int i = 0; i < data.size(); i++)
if (data[i] <= kth_element)
result.push_back({i, data[i]});
return result;
}
update: It is also possible to find the kth_element by maintaing a max heap with size at most k.
It only need O(k) memory for heap instead of O(n) memory in the nth_element method.
It needs O(n log k) time but if k is small then i think it should be comparable to O(n) method.
I am not sure about it but my reason are the heap may be cached and you don't need to spend time for copying data.
vector<idx_pair> getKSmallest(vector<double> const& data, int k)
{
priority_queue<double> pq;
for (auto d : data){
if (pq.size() >= k && pq.top() > d){
pq.push(d)
pq.pop();
}
else if (pq.size() < k)
pq.push(d);
}
double kth_element = pq.top();
vector<idx_pair> result;
for (int i = 0; i < data.size(); i++)
if (data[i] <= kth_element)
result.push_back({i, data[i]});
return result;
}