Language: C++

#define debug_interval ','
#define dump_interval ' '
#define debug_toggle 1

//{
#pragma GCC optimize ("O3","unroll-loops")
#pragma GCC target ("avx")
#include<bits/stdc++.h>
using namespace std;
#define hi cerr<<"hi"<<endl;
#define int long long
#define INT_MAX LLONG_MAX
#define rep(i,n) for(int i=0;i<(n);i++)
#define rep1(i,n) for(int i=1;i<=(n);i++)
#define pb push_back
#define all(v) (v).begin(),(v).end()
#define rall(v) (v).rbegin(),(v).rend()
#define fi first
#define se second
#define mp make_pair
#define rev reverse
#define dans dump(ans)
#define MOD 1000000007
#define amp(v,n) (v).count(n)?v[n]++:v[n]=1
#define sysp system("pause")
#define PI acos(-1)
#define ins insert
#define aer(v) (v).erase((v).begin(),(v).end()) 

//{
inline int toInt(string s){int v;istringstream sin(s);sin>>v;return v;}
template<class T> inline string toString(T x){ostringstream sout;sout<<x;return sout.str();}
template<class...A> inline void dump(){cout<<endl;}
template<class...A> inline void dump_rest() {cout<<endl;}
template<class T, class...A> inline void dump_rest(const T& first, const A&... rest){cout<<dump_interval<<first;dump_rest(rest...);}
template<class T,class...A> inline void dump(const T&first,const A&...rest){cout<<first;dump_rest(rest...);}
template<class...A> inline void debug(){cerr<<endl;}
template<class...A> inline void debug_rest() {cerr<<endl;}
template<class T, class...A> inline void debug_rest(const T& first, const A&... rest){cerr<<debug_interval<<first;debug_rest(rest...);}
template<class T,class...A> inline void debug(const T&first,const A&...rest){if(debug_toggle)cerr<<first,debug_rest(rest...);}
//}

typedef vector<int>vint;
typedef pair<int,int>pint;
typedef vector<pint>vpint;

template<typename A,typename B>inline void chmin(A &a,B b){if(a>b)a=b;}
template<typename A,typename B>inline void chmax(A &a,B b){if(a<b)a=b;}
template<int m>class mint{private:int i;public:mint() : i(0){}mint(int i): i((i%m+m)%m){}mint operator+(const mint& o){return o.i+i;}mint operator*(const mint& o){return o.i*i;}mint operator-(){return -i;}operator int() {return i;}};

//}

int n,k;
vector<priority_queue<int> > a[111111];
priority_queue<pint> pq;

main(){
	cin>>n>>k;
	int t[n],d[n],used[n]={},ans=0,dc[n]={};
	rep(i,n)cin>>t[i]>>d[i];
	rep(i,n){
		pq.push({d[i],t[i]});
		a[d[i]].push(t[i]);
	}
	int cnt=0;
	priority_queue<pint,vector<pint>,greater<pint> > q;
	rep(i,k){
		pint x=pq.top();pq.pop();
		used[x.se]=1;
		dc[x.se]++;
		ans+=x.fi;
		//v.pb(x.fi);
		q.push(x);
	}
	int div=0;
	bool flg=0;
	rep(i,n)if(used[i])div++;
	ans+=div*div;
	for(int i=0;i<k-div&&!pq.empty();){
		if(used[pq.top().se]){
			pq.pop();
			continue;
		}
		else {
			pint tmp=q.top();q.pop();
			dc[tmp.se]--;
			if(dc[tmp.se]==0)flg=1;
		}
	}
	dans;
}


Compiler output:
a.cc:12:9: warning: "INT_MAX" redefined
   12 | #define INT_MAX LLONG_MAX
      |         ^~~~~~~
In file included from /usr/include/c++/14/climits:42,
                 from /usr/include/x86_64-linux-gnu/c++/14/bits/stdc++.h:38,
                 from a.cc:8:
/usr/lib/gcc/x86_64-linux-gnu/14/include/limits.h:120:9: note: this is the location of the previous definition
  120 | #define INT_MAX __INT_MAX__
      |         ^~~~~~~
a.cc:57:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
   57 | main(){
      | ^~~~
a.cc: In function 'int main()':
a.cc:63:25: error: 'class std::vector<std::priority_queue<long long int> >' has no member named 'push'
   63 |                 a[d[i]].push(t[i]);
      |                         ^~~~