Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;

// --
#define HOGE                                                                   \
  std::cin.tie(0);                                                             \
  std::ios::sync_with_stdio(false);
template <typename T> void CIN_VAR(T &t) { cin >> t; }
template <typename A, typename... R> void CIN_VAR(A &car, R &...cdr) {
  std::cin >> car;
  CIN_VAR(cdr...);
}
#define VAR(type, ...)                                                         \
  type __VA_ARGS__;                                                            \
  CIN_VAR(__VA_ARGS__);
#define VEC(t, c, n)                                                           \
  std::vector<t> c(n);                                                         \
  for (auto &i : c)                                                            \
    std::cin >> i;
#define MAT(t, c, m, n)                                                        \
  std::vector<std::vector<t>> c(m, std::vector<t>(n));                         \
  for (auto &row : c)                                                          \
    for (auto &w : row)                                                        \
      std::cin >> w;

#define ALL(c) (c).begin(), (c).end()
#define CALL(c) (c).cbegin(), (c).cend()
#define RALL(c) (c).rbegin(), (c).rend()
#define CRALL(c) (c).crbegin(), (c).crend()
#define REP(i, n) for (int i = 0; i < static_cast<int>(n); i++)
#define FOR(i, s, n) for (int i = s; i < static_cast<int>(n); i++)
#define PSORT(s, p, e) partial_sort((s), next((s), (p)), (e))

#define IN(a, x, b) (a <= x && x < b)
template <class T> inline T CHMAX(T &x, const T y) {
  return x = (x < y) ? y : x;
}
template <class T> inline T CHMIN(T &x, const T y) {
  return x = (x > y) ? y : x;
}
#define SDESC(t) [](t &x, t &y) { return x > y; }

#define VIEW(v, t)                                                             \
  for_each((v).begin(), (v).end(), [](t &a) { cout << a << " "; });            \
  cout << "\n";
#define COUT(dist) std::cout << (dist);
#define WS std::cout << " ";
#define NL std::cout << "\n";
// --

// --
using PII = pair<int, int>;
using PLL = pair<long long, long long>;
using LL = long long;
using ULL = unsigned long long;

template <class T> using V = vector<T>;
template <class T> using V2 = vector<vector<T>>;
template <class T, int W = 255> using A = array<T, W>;
template <class T, int W = 255, int H = 255> using A2 = array<array<T, W>, W>;

constexpr int dx[] = {-1, 0, 1, 0};
constexpr int dy[] = {0, -1, 0, 1};
// --

LL n, k;
V<LL> t, d;
V<PLL> pv;
std::set<int> s;
std::stack<int> st;

int main() {
  HOGE;
  cin >> n;
  cin >> k;
  t.resize(n);
  d.resize(n);
  pv.resize(n);
  REP(i, n) {
    cin >> t[i];
    cin >> d[i];
  }
  {
    REP(i, n) pv[i] = PLL(d[i], t[i]);
    std::sort(ALL(pv), SDESC(PLL));
    REP(i, n) std::tie(d[i], t[i]) = pv[i];
  }

  LL ans = 0;
  REP(i, k) {
    if (s.count(t[i]))
      st.push(d[i]);
    s.insert(t[i]);
    ans += d[i];
  }
  LL sum = ans;
  ans += (LL)((LL)s.size() * (LL)s.size());
  FOR(i, k, n) {
    if (s.empty())
      break;
    if (s.count(t[i]))
      continue;
    s.insert(t[i]);
    sum -= st.top();
    st.pop();
    sum += d[i];
    CHMAX<LL>(ans, sum + (LL)(s.size() * s.size()));
  }
  COUT(ans) NL;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;

// --
#define HOGE                                                                   \
  std::cin.tie(0);                                                             \
  std::ios::sync_with_stdio(false);
template <typename T> void CIN_VAR(T &t) { cin >> t; }
template <typename A, typename... R> void CIN_VAR(A &car, R &...cdr) {
  std::cin >> car;
  CIN_VAR(cdr...);
}
#define VAR(type, ...)                                                         \
  type __VA_ARGS__;                                                            \
  CIN_VAR(__VA_ARGS__);
#define VEC(t, c, n)                                                           \
  std::vector<t> c(n);                                                         \
  for (auto &i : c)                                                            \
    std::cin >> i;
#define MAT(t, c, m, n)                                                        \
  std::vector<std::vector<t>> c(m, std::vector<t>(n));                         \
  for (auto &row : c)                                                          \
    for (auto &w : row)                                                        \
      std::cin >> w;

#define ALL(c) (c).begin(), (c).end()
#define CALL(c) (c).cbegin(), (c).cend()
#define RALL(c) (c).rbegin(), (c).rend()
#define CRALL(c) (c).crbegin(), (c).crend()
#define REP(i, n) for (int i = 0; i < static_cast<int>(n); i++)
#define FOR(i, s, n) for (int i = s; i < static_cast<int>(n); i++)
#define PSORT(s, p, e) partial_sort((s), next((s), (p)), (e))

#define IN(a, x, b) (a <= x && x < b)
template <class T> inline T CHMAX(T &x, const T y) {
  return x = (x < y) ? y : x;
}
template <class T> inline T CHMIN(T &x, const T y) {
  return x = (x > y) ? y : x;
}
#define SDESC(t) [](t &x, t &y) { return x > y; }

#define VIEW(v, t)                                                             \
  for_each((v).begin(), (v).end(), [](t &a) { cout << a << " "; });            \
  cout << "\n";
#define COUT(dist) std::cout << (dist);
#define WS std::cout << " ";
#define NL std::cout << "\n";
// --

// --
using PII = pair<int, int>;
using PLL = pair<long long, long long>;
using LL = long long;
using ULL = unsigned long long;

template <class T> using V = vector<T>;
template <class T> using V2 = vector<vector<T>>;
template <class T, int W = 255> using A = array<T, W>;
template <class T, int W = 255, int H = 255> using A2 = array<array<T, W>, W>;

constexpr int dx[] = {-1, 0, 1, 0};
constexpr int dy[] = {0, -1, 0, 1};
// --

LL n, k;
V<LL> t, d;
V<PLL> pv;
std::set<int> s;
std::stack<int> st;

int main() {
  HOGE;
  cin >> n;
  cin >> k;
  t.resize(n);
  d.resize(n);
  pv.resize(n);
  REP(i, n) {
    cin >> t[i];
    cin >> d[i];
  }
  {
    REP(i, n) pv[i] = PLL(d[i], t[i]);
    std::sort(ALL(pv), SDESC(PLL));
    REP(i, n) std::tie(d[i], t[i]) = pv[i];
  }

  LL ans = 0;
  REP(i, k) {
    if (s.count(t[i]))
      st.push(d[i]);
    s.insert(t[i]);
    ans += d[i];
  }
  LL sum = ans;
  ans += (LL)((LL)s.size() * (LL)s.size());
  FOR(i, k, n) {
    if (st.empty())
      break;
    if (s.count(t[i]))
      continue;
    s.insert(t[i]);
    sum -= st.top();
    st.pop();
    sum += d[i];
    CHMAX<LL>(ans, sum + (LL)(s.size() * s.size()));
  }
  COUT(ans) NL;
}