fork download
  1. #include <iostream>
  2.  
  3. #include <vector>
  4. #include <queue>
  5. #include <chrono>
  6. #include <limits>
  7. #include <algorithm>
  8. #include<cmath>
  9.  
  10. using namespace std::chrono;
  11.  
  12. typedef std::pair<double, int> idx_pair;
  13. typedef std::priority_queue<idx_pair, std::vector<idx_pair>, std::greater<idx_pair>> idx_queue;
  14.  
  15. std::vector<idx_pair> getKSmallest(std::vector<double> const& data, int k)
  16. {
  17. idx_queue q;
  18. {
  19. std::vector<idx_pair> idxPairs(data.size());
  20. for (auto i = 0; i < data.size(); i++)
  21. idxPairs[i] = idx_pair(data[i], i);
  22.  
  23. q = idx_queue(std::begin(idxPairs), std::end(idxPairs));
  24. };
  25.  
  26. std::vector<idx_pair> result;
  27.  
  28. auto topPop = [&q, &result]()
  29. {
  30. result.push_back(q.top());
  31. q.pop();
  32. };
  33.  
  34. for (auto i = 0; i < k; i++)
  35. topPop();
  36.  
  37. auto const largest = result.back().first;
  38. while (q.empty() == false)
  39. {
  40. if (q.top().first == largest)
  41. topPop();
  42. else
  43. break;
  44. }
  45.  
  46. return result;
  47. }
  48.  
  49. void printVec(std::vector<idx_pair> const& data)
  50. {
  51. for (auto const d : data)
  52. std::cout << "{" << d.second << "," << d.first << "} ";
  53.  
  54. std::cout << std::endl;
  55. }
  56.  
  57. bool equal(double value1, double value2)
  58. {
  59. return value1 == value2 || std::abs(value2 - value1) <= std::numeric_limits<double>::epsilon();
  60. }
  61.  
  62. std::vector<idx_pair> getNSmallest(std::vector<double> const& data, int n)
  63. {
  64. std::vector<idx_pair> idxPairs(data.size());
  65. for (auto i = 0; i < data.size(); i++)
  66. idxPairs[i] = idx_pair(data[i], i);
  67.  
  68. std::nth_element(std::begin(idxPairs), std::begin(idxPairs) + n, std::end(idxPairs));
  69.  
  70. std::vector<idx_pair> result(std::begin(idxPairs), std::begin(idxPairs) + n);
  71.  
  72. auto const largest = result.back().first;
  73. for (auto it = std::begin(idxPairs) + n; it != std::end(idxPairs); ++it)
  74. {
  75. if (equal(it->first, largest))
  76. result.push_back(*it);
  77. }
  78.  
  79. return result;
  80. }
  81.  
  82. int main() {
  83. std::vector<double> data = { 3.3, 1.1, 6.5, 4.2, 1.1, 3.5 };
  84.  
  85. auto t1 = high_resolution_clock::now();
  86. auto const result = getNSmallest(data, 3);
  87. auto t2 = high_resolution_clock::now();
  88.  
  89. std::cout << "time: " << duration_cast<milliseconds>(t2 - t1).count() << std::endl;
  90.  
  91. printVec(result);
  92.  
  93. return 0;
  94. }
Success #stdin #stdout 0s 3476KB
stdin
Standard input is empty
stdout
time: 0
{4,1.1} {1,1.1} {0,3.3}