Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <array>
#include <assert.h>
#include <bitset>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <stdio.h>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

#define pb push_back
#define fi first
#define se second
#define all(x) (x).begin(), (x).end()
#define mn(x) *min_element((x).begin(), (x).end())
#define mx(x) *max_element((x).begin(), (x).end())
#define acc(x) accumulate((x).begin(), (x).end(), 0ll)
#define eb emplace_back
#define el '\n'
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
typedef pair<ll, ll> pll;
typedef pair<int, int> pii;
typedef vector<ll> vll;
typedef vector<int> vi;
typedef vector<pii> vpi;
typedef vector<pll> vpll;
typedef vector<vi> vvi;
typedef vector<vll> vvll;
typedef vector<bool> vb;
typedef vector<string> vs;

const ld pi = acos(-1);
const ll inf = (ll)1e12 + 1;
// const int mod=1000000007;
const int mod = 998244353;

int n, m, k;
const int N = 3e5 + 10;

struct Set_with_sum {
  struct node {
    int no;
    ll sum;
    node *l, *r;
    node() : no(0), sum(0ll), l(NULL), r(NULL) {}
  };
  node *root;
  int lo = 0, hi;
  Set_with_sum(int hii) : root(new node), hi(hii){};

  void insert(node *p, int l, int r, int id) {
    if (l == r) {
      p->sum += id;
      p->no++;
      return;
    }
    int mid = (l + r) >> 1;
    if (p->l == NULL)
      p->l = new node;
    if (p->r == NULL)
      p->r = new node;
    if (id > mid) {
      insert(p->r, mid + 1, r, id);
    } else {
      insert(p->l, l, mid, id);
    }
    p->sum = p->l->sum + p->r->sum;
    p->no = p->l->no + p->r->no;
  }

  void erase(node *p, int l, int r, int id) {
    if (l == r) {
      p->sum -= id;
      p->no--;
      return;
    }
    int mid = (l + r) >> 1;
    if (id > mid) {
      erase(p->r, mid + 1, r, id);
    } else {
      erase(p->l, l, mid, id);
    }
    p->sum = p->l->sum + p->r->sum;
    p->no = p->l->no + p->r->no;
  }

  ll query(node *p, int l, int r, int k) {
    if (!k)
      return 0;
    if (l == r)
      return (ll)k * l;
    if (p->l == NULL)
      p->l = new node;
    if (p->r == NULL)
      p->r = new node;
    int m = (l + r) >> 1;
    if (p->r->no > k) {
      return query(p->r, m + 1, r, k);
    } else {
      return p->r->sum + query(p->l, l, m, k - p->r->no);
    }
  }
  ll query(int k) { return query(root, 0, hi, k); }
  void insert(int no) { insert(root, 0, hi, no); }
  void erase(int no) { erase(root, 0, hi, no); }
};

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cin >> k >> n;
  Set_with_sum s(1e9 + 1);
  priority_queue<pii, vpi> pq;
  for (int i = 0, u, v; i < n; i++) {
    cin >> u >> v;
    s.insert(v);
    pq.emplace(v, u);
  }
  vb mark(n + 1);
  ll diff = 0;
  ll a = 0, already = 0;
  while (!pq.empty()) {
    auto f = pq.top();
    pq.pop();
    if (!mark[f.se]) {
      mark[f.se] = true;
      already += f.fi;
      diff++;
      s.erase(f.fi);
      a = max(a, already + s.query(k - (int)diff) + diff * diff);
      if (diff == k)
        break;
    }
  }

  cout << a << el;

  return 0;
}


Fixed submission:
#include <algorithm>
#include <array>
#include <assert.h>
#include <bitset>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <stdio.h>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

#define pb push_back
#define fi first
#define se second
#define all(x) (x).begin(), (x).end()
#define mn(x) *min_element((x).begin(), (x).end())
#define mx(x) *max_element((x).begin(), (x).end())
#define acc(x) accumulate((x).begin(), (x).end(), 0ll)
#define eb emplace_back
#define el '\n'
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
typedef pair<ll, ll> pll;
typedef pair<int, int> pii;
typedef vector<ll> vll;
typedef vector<int> vi;
typedef vector<pii> vpi;
typedef vector<pll> vpll;
typedef vector<vi> vvi;
typedef vector<vll> vvll;
typedef vector<bool> vb;
typedef vector<string> vs;

const ld pi = acos(-1);
const ll inf = (ll)1e12 + 1;
// const int mod=1000000007;
const int mod = 998244353;

int n, m, k;
const int N = 3e5 + 10;

struct Set_with_sum {
  struct node {
    int no;
    ll sum;
    node *l, *r;
    node() : no(0), sum(0ll), l(NULL), r(NULL) {}
  };
  node *root;
  int lo = 0, hi;
  Set_with_sum(int hii) : root(new node), hi(hii){};

  void insert(node *p, int l, int r, int id) {
    if (l == r) {
      p->sum += id;
      p->no++;
      return;
    }
    int mid = (l + r) >> 1;
    if (p->l == NULL)
      p->l = new node;
    if (p->r == NULL)
      p->r = new node;
    if (id > mid) {
      insert(p->r, mid + 1, r, id);
    } else {
      insert(p->l, l, mid, id);
    }
    p->sum = p->l->sum + p->r->sum;
    p->no = p->l->no + p->r->no;
  }

  void erase(node *p, int l, int r, int id) {
    if (l == r) {
      p->sum -= id;
      p->no--;
      return;
    }
    int mid = (l + r) >> 1;
    if (id > mid) {
      erase(p->r, mid + 1, r, id);
    } else {
      erase(p->l, l, mid, id);
    }
    p->sum = p->l->sum + p->r->sum;
    p->no = p->l->no + p->r->no;
  }

  ll query(node *p, int l, int r, int k) {
    if (!k)
      return 0;
    if (l == r)
      return (ll)k * l;
    if (p->l == NULL)
      p->l = new node;
    if (p->r == NULL)
      p->r = new node;
    int m = (l + r) >> 1;
    if (p->r->no > k) {
      return query(p->r, m + 1, r, k);
    } else {
      return p->r->sum + query(p->l, l, m, k - p->r->no);
    }
  }
  ll query(int k) { return query(root, 0, hi, k); }
  void insert(int no) { insert(root, 0, hi, no); }
  void erase(int no) { erase(root, 0, hi, no); }
};

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cin >> n >> k;
  Set_with_sum s(1e9 + 1);
  priority_queue<pii, vpi> pq;
  for (int i = 0, u, v; i < n; i++) {
    cin >> u >> v;
    s.insert(v);
    pq.emplace(v, u);
  }
  vb mark(n + 1);
  ll diff = 0;
  ll a = 0, already = 0;
  while (!pq.empty()) {
    auto f = pq.top();
    pq.pop();
    if (!mark[f.se]) {
      mark[f.se] = true;
      already += f.fi;
      diff++;
      s.erase(f.fi);
      a = max(a, already + s.query(k - (int)diff) + diff * diff);
      if (diff == k)
        break;
    }
  }

  cout << a << el;

  return 0;
}