999ms

Untitled

Oct 17th, 2021
953
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #pragma GCC Optimize("Ofast")
  2.  
  3. #include <bits/stdc++.h>
  4.  
  5. #define all(x) begin(x),end(x)
  6.  
  7. using namespace std;
  8. using ld = long double;
  9. using ll = long long;
  10.  
  11. const bool LOCAL = true;
  12.  
  13. void solve() {
  14.   int n, k;
  15.   cin >> n >> k;
  16.   vector<pair<int, int>> arr(n);
  17.   map<int, map<int, int>> weightToCountedColors; // weight -> {color -> count}
  18.   for (int i = 0; i < n; i++) {
  19.     cin >> arr[i].first >> arr[i].second; // weight, color
  20.     weightToCountedColors[arr[i].first][arr[i].second]++;
  21.   }
  22.  
  23.   // block = a set of elements with the same weight
  24.   // multiblock = block with elements with at least two different colors
  25.   // simpleblock = block with elements with only one color (<=> number of elements in the block is one)
  26.  
  27.   vector<pair<int, map<int, int>>> blocks;
  28.  
  29.   for (auto &element: weightToCountedColors) {
  30.     auto &weight = element.first;
  31.     auto &block = element.second;
  32.  
  33.     if (block.size() >= 2) {
  34.       blocks.emplace_back(weight, block); // current block is multiblock
  35.       continue;
  36.     }
  37.  
  38.     block.begin()->second = 1; // current block is a simpleblock, so we need to set number of elements to 1
  39.     int currentColor = block.begin()->first;
  40.  
  41.     if (blocks.empty() || blocks.back().second.size() >= 2) {
  42.       blocks.emplace_back(weight, block); // previous block is a multiblock or it is first element
  43.       continue;
  44.     }
  45.  
  46.     if (blocks.back().second.begin()->first == currentColor) {
  47.       blocks.back().first = weight; // previous block is a simpleblock with the same color as current
  48.     } else {
  49.       blocks.emplace_back(weight, block); // previous block is a simpleblock with another color
  50.     }
  51.   }
  52.  
  53.   // The size of answer must be equal to at most k
  54.  
  55.   long long answer = 0;
  56.   int currentSize = 0;
  57.   for (auto it = blocks.rbegin(); it != blocks.rend(); it++) {
  58.     auto &weight = it->first;
  59.     auto &block = it->second;
  60.  
  61.     for (auto element: block) {
  62.       auto count = element.second;
  63.       currentSize += count;
  64.       answer += 1ll * count * weight;
  65.     }
  66.  
  67.     if (currentSize >= k) {
  68.       while (currentSize > k) {
  69.         currentSize--;
  70.         answer -= weight;
  71.       }
  72.       break;
  73.     }
  74.   }
  75.  
  76.   cout << answer << '\n';
  77. }
  78.  
  79. int main() {
  80.   ios_base::sync_with_stdio(false);
  81.   cin.tie(nullptr);
  82.   cout.tie(nullptr);
  83.   if (LOCAL) {
  84.     freopen("input.txt", "r", stdin);
  85.     freopen("output.txt", "w", stdout);
  86.   }
  87.  
  88.   solve();
  89. }
RAW Paste Data