fork download
  1. /*******************************************************
  2. 3) Amazon OA — k-th Smallest in Every Window of Length m
  3.  
  4. Problem:
  5. - For each sliding window of size m, output the k-th smallest element.
  6.  
  7. Idea:
  8. - Keep two multisets:
  9.   small = k smallest numbers in window
  10.   big = remaining numbers
  11. - Then k-th smallest = max(small)
  12.  
  13. Time: O(n log m)
  14. *******************************************************/
  15.  
  16. #include <bits/stdc++.h> // standard CP header
  17. using namespace std; // use std names directly
  18.  
  19. static void rebalance(multiset<long long>& small, multiset<long long>& big, int k) { // keep sets valid
  20. while ((int)small.size() > k) { // if small has too many
  21. auto it = prev(small.end()); // iterator to largest in small
  22. big.insert(*it); // move it to big
  23. small.erase(it); // remove from small
  24. }
  25. while ((int)small.size() < k && !big.empty()) { // if small has too few and big has elements
  26. auto it = big.begin(); // iterator to smallest in big
  27. small.insert(*it); // move it to small
  28. big.erase(it); // remove from big
  29. }
  30.  
  31. while (!small.empty() && !big.empty()) { // if both sets non-empty
  32. auto itSmallMax = prev(small.end()); // max element in small
  33. auto itBigMin = big.begin(); // min element in big
  34. if (*itSmallMax <= *itBigMin) break; // order is fine, stop
  35. long long a = *itSmallMax; // value from small that is too large
  36. long long b = *itBigMin; // value from big that is too small
  37. small.erase(itSmallMax); // remove a
  38. big.erase(itBigMin); // remove b
  39. small.insert(b); // put b into small
  40. big.insert(a); // put a into big
  41. }
  42. }
  43.  
  44. int main() { // program start
  45. ios::sync_with_stdio(false); // fast I/O
  46. cin.tie(nullptr); // fast I/O
  47.  
  48. int n, m, k; // n=array size, m=window size, k=order statistic
  49. cin >> n >> m >> k; // read them
  50. vector<long long> a(n); // array values
  51. for (int i = 0; i < n; i++) cin >> a[i]; // read array
  52.  
  53. multiset<long long> small, big; // the two multisets
  54.  
  55. for (int i = 0; i < m; i++) { // insert first window
  56. big.insert(a[i]); // put everything in big first
  57. }
  58. rebalance(small, big, k); // then move smallest k into small properly
  59.  
  60. auto print_ans = [&]() { // helper lambda to print k-th smallest
  61. cout << *prev(small.end()); // k-th smallest is maximum of small
  62. };
  63.  
  64. print_ans(); // print answer for first window
  65.  
  66. for (int i = m; i < n; i++) { // slide window from i=m to n-1
  67. long long out = a[i - m]; // element leaving the window
  68. long long in = a[i]; // element entering the window
  69.  
  70. auto itSmall = small.find(out); // try to find outgoing element in small
  71. if (itSmall != small.end()) { // if found in small
  72. small.erase(itSmall); // erase it
  73. } else { // else it must be in big
  74. auto itBig = big.find(out); // find in big
  75. if (itBig != big.end()) big.erase(itBig); // erase it if found
  76. }
  77.  
  78. if (!small.empty() && in <= *prev(small.end())) { // if incoming likely belongs in small
  79. small.insert(in); // insert into small
  80. } else { // otherwise
  81. big.insert(in); // insert into big
  82. }
  83.  
  84. rebalance(small, big, k); // restore invariants
  85.  
  86. cout << " "; // space between answers
  87. print_ans(); // print current window answer
  88. }
  89.  
  90. cout << "\n"; // newline at end
  91. return 0; // done
  92. }
Runtime error #stdin #stdout 0.04s 5316KB
stdin
Standard input is empty
stdout
Standard output is empty