kmjp's blog

競技プログラミング参加記です

Codeforces #635 Div1 B. Xenia and Colorful Gems

どんどん行こう。
https://codeforces.com/contest/1336/problem/B

問題

3色の宝石群があり、個々の宝石の重さが与えられる。
各色1個ずつ宝石を選び、宝石対の重さの差の2乗和を最小化したい。
その値を答えよ。

解法

もし3つのうち最大の重さと最小の重さが決まっているとき、残り1色はその中間に近い重さが良い。
これをシミュレートしていこう。

全宝石を重さ順にソートし、各色の列に1個ずつ追加していこう。
そのたび、3色中2色で最大の重さの宝石を選んだとして、残り1色の最適な重さを二分探索で求める、ということを2色の選び方を総当たりしつつ行えばよい。

int T;
int N[3];
int A[3][101010];
vector<pair<int,int>> V;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	scanf("%d",&T);
	while(T--) {
		V.clear();
		scanf("%d%d%d",&N[0],&N[1],&N[2]);
		FOR(i,3) {
			FOR(j,N[i]) {
				scanf("%d",&x);
				A[i][j]=x;
				V.push_back({x,i});
			}
			sort(A[i],A[i]+N[i]);
		}
		sort(ALL(V));
		ll mi=1LL<<62;
		ll t[3]={-1LL<<30,-1LL<<30,-1LL<<30};
		FORR(e,V) {
			t[e.second]=e.first;
			if(t[0]<0 || t[1]<0 || t[2]<0) continue;
			FOR(x,3) FOR(y,3) if(x!=y) {
				ll tar=(t[x]+t[y])/2;
				i=3-x-y;
				j=lower_bound(A[i],A[i]+N[i],tar)-A[i];
				for(k=max(0,j-1);k<=min(N[i]-1,j+1);k++) {
					ll v=A[i][k];
					mi=min(mi,(t[x]-t[y])*(t[x]-t[y])+(t[x]-v)*(t[x]-v)+(t[y]-v)*(t[y]-v));
				}
			}
		}
		cout<<mi<<endl;
		
		
		
	}
}

まとめ

本番これ割と躊躇なく通してるな。今考えてもあってるのか少し悩んでしまった。