ABC123 D - Cake 123 (400)

D - Cake 123

概要

X項の数列A、Y項の数列B、Z項の数列Cがある。
A、B、Cからそれぞれ1項ずつ取り出して和を求めるとき、その和が大きい順にK個出力せよ。

制約

1 ≦ X, Y, Z ≦ 1000
1 ≦ K ≦ min(3000, X * Y * Z)
1 ≦ Ai, Bi, Ci ≦ 10^10

方針

Ai + Bj + Ckの組み合わせをすべて試すと10^9の計算量で間に合わない。

半分全列挙的な考え方を使う。
まず、Ai + Bjのすべての和を算出し、配列vに格納してソートする。
その後、viとCjについて、それぞれ上位K個同士で和を算出し、配列wに格納してソートする。
こうすると、計算量も10^7程度で間に合うし、最終的に使用するk個を確実に見つけることができる。
あとはwの上位k個を出力すればよい。

解答

Submission #4859444 - AtCoder Beginner Contest 123

#include <bits/stdc++.h>
#define rep(i,n) for(int i=(0);i<(n);i++)
 
using namespace std;
 
typedef long long ll;
 
int main()
{
	cin.tie(0);
	ios::sync_with_stdio(false);
 
	int X, Y, Z, K;
	cin >> X >> Y >> Z >> K;
 
	vector<ll> A(X);
	rep(i, X) cin >> A[i];
 
	vector<ll> B(X);
	rep(i, Y) cin >> B[i];
 
	vector<ll> C(X);
	rep(i, Z) cin >> C[i];
 
	vector<ll> v;
	for(int i = 0; i < X; i++){
		for(int j = 0; j < Y; j++){
			v.push_back(A[i] + B[j]);
		}
	}
 
	sort(v.begin(), v.end(), greater<ll>());
 
	vector<ll> w;
	for(int i = 0; i < min(K, X * Y); i++){
		for(int j = 0; j < min(K, Z); j++){
			w.push_back(v[i] + C[j]);
		}
	}
 
	sort(w.begin(), w.end(), greater<ll>());
	rep(i, K) cout << w[i] << endl;
}