Language: C++

#include <array>
#include <bitset>
#include <algorithm>

using namespace std;

struct sushi_t {
  int d;
  int t;
};

static int n, k;
static std::array<sushi_t, 100000> sushis{};
static std::array<int, 100000> takenNetas{};
static std::bitset<100000> takenSushis{};

int64_t solve() {
  int64_t ret0 = 0;
  int64_t nn = 0;
  for (int i = 0; i < k; ++i) {
    const auto& s = sushis[i];
    ret0 += int64_t(s.d);
    if (takenNetas[s.t] == 0) {
      nn++;
    }
    takenNetas[s.t]++;
    takenSushis.set(i);
  }
  int64_t ret = ret0 + nn * nn;

  for (++nn; nn <= k; ++nn) {
    bool found = false;
    for (int i = n - 1; i >= 0; --i) {
      if (!takenSushis.test(i)) {
        continue;
      }
      const auto& s = sushis[i];
      assert(takenNetas[s.t] >= 1);
      if (takenNetas[s.t] <= 1) {
        continue;
      }
      ret0 -= int64_t(s.d);
      takenNetas[s.t]--;
      takenSushis.reset(i);
      found = true;
      break;
    }
    if (!found) {
      break;
    }
    found = false;
    for (int i = 0; i < n; ++i) {
      if (takenSushis.test(i)) {
        continue;
      }
      const auto& s = sushis[i];
      if (takenNetas[s.t] > 0) {
        continue;
      }
      ret0 += int64_t(s.d);
      takenNetas[s.t]++;
      takenSushis.set(i);
      ret = std::max(ret0 + nn * nn, ret);
      found = true;
      break;
    }
    if (!found) {
      break;
    }
  }
  return ret;
}

int main() {
  std::cin.tie(nullptr);
  std::ios::sync_with_stdio(false);
  std::cin >> n >> k;
  for (int i = 0; i < n; ++i) {
    auto& s = sushis[i];
    std::cin >> s.t >> s.d;
    s.t--;
  }
  std::sort(sushis.begin(), sushis.begin() + n,
            [](const auto& lhs, const auto& rhs) { return lhs.d > rhs.d; });
  std::cout << solve() << std::endl;
  return 0;
}

Compiler output:
a.cc: In function 'int64_t solve()':
a.cc:38:7: error: 'assert' was not declared in this scope
   38 |       assert(takenNetas[s.t] >= 1);
      |       ^~~~~~
a.cc:4:1: note: 'assert' is defined in header '<cassert>'; this is probably fixable by adding '#include <cassert>'
    3 | #include <algorithm>
  +++ |+#include <cassert>
    4 | 
a.cc: In function 'int main()':
a.cc:75:8: error: 'cin' is not a member of 'std'
   75 |   std::cin.tie(nullptr);
      |        ^~~
a.cc:4:1: note: 'std::cin' is defined in header '<iostream>'; this is probably fixable by adding '#include <iostream>'
    3 | #include <algorithm>
  +++ |+#include <iostream>
    4 | 
a.cc:76:13: error: incomplete type 'std::ios' {aka 'std::basic_ios<char>'} used in nested name specifier
   76 |   std::ios::sync_with_stdio(false);
      |             ^~~~~~~~~~~~~~~
a.cc:77:8: error: 'cin' is not a member of 'std'
   77 |   std::cin >> n >> k;
      |        ^~~
a.cc:77:8: note: 'std::cin' is defined in header '<iostream>'; this is probably fixable by adding '#include <iostream>'
a.cc:80:10: error: 'cin' is not a member of 'std'
   80 |     std::cin >> s.t >> s.d;
      |          ^~~
a.cc:80:10: note: 'std::cin' is defined in header '<iostream>'; this is probably fixable by adding '#include <iostream>'
a.cc:85:8: error: 'cout' is not a member of 'std'
   85 |   std::cout << solve() << std::endl;
      |        ^~~~
a.cc:85:8: note: 'std::cout' is defined in header '<iostream>'; this is probably fixable by adding '#include <iostream>'
a.cc:85:32: error: 'endl' is not a member of 'std'
   85 |   std::cout << solve() << std::endl;
      |                                ^~~~
a.cc:4:1: note: 'std::endl' is defined in header '<ostream>'; this is probably fixable by adding '#include <ostream>'
    3 | #include <algorithm>
  +++ |+#include <ostream>
    4 |