priority_queue で動的に $K$ 個出力する問題作りたいな〜と思ってぼんやりストックしてた問題と一緒だった！！！！！！

問題概要

$X$ 個の要素からなる数列 $A_1, \dots, A_X$
$Y$ 個の要素からなる数列 $B_1, \dots, B_Y$
$Z$ 個の要素からなる数列 $C_1, \dots, C_Z$

が与えられる。 $A, B, C$ からそれぞれ 1 個ずつとりだして和をとった値は $XYZ$ 個考えられるが、それらを大きい順に $K$ 個出力せよ。

制約

$1 \le K \le 3000$
$1 \le X, Y, Z \le 1000$

考えたこと (editorial 解法 3)

$A, B, C$ はそれぞれ大きい順にソートしておく。まず一番大きいのは $A, B, C$ それぞれの最大値をとったもの $A_1 + B_1 + C_1$ で決まりである。ここから priority_queue をうまく用いる。

この次に大きいものは

$A_2 + B_1 + C_1$
$A_1 + B_2 + C_1$
$A_1 + B_1 + C_2$

の 3 つの候補があるが、それぞれをキューに push しておく。このときにキューの先頭にあるものが、2 番目に大きい値になる。

同様にその 2 番目に大きい値を pop して来て、それに対し

A のみ index を 1 個進めて足したもの
B のみ index を 1 個進めて足したもの
C のみ index を 1 個進めて足したもの

の 3 通りの値をキューに push していく。そしてこの時点でキューの先頭にあるものが 3 番目の値になる。以後これを繰り返して行く。注意点として、「一度 push したものは二度 push しない」という管理が必要となる。

計算量は、

キューから pop する回数が $K$ 回
その $K$ 回それぞれについて、キューの push するのは 3 回以下なので、キューに push する回数も多くても $3K$ 以下

ということで、計算量は $O(K\log{K})$ となる。

#include <iostream>
#include <vector>
#include <algorithm>
#include <queue>
#include <set>
using namespace std;
using Data = pair<long long, vector<int> >;

int main() {
    vector<long long> N(3);
    int K;
    for (int i = 0; i < 3; ++i) cin >> N[i]; cin >> K;
    vector<vector<long long> > v(3);
    for (int iter = 0; iter < 3; ++iter) {
        v[iter].resize(N[iter]);
        for (int i = 0; i < N[iter]; ++i) cin >> v[iter][i];
        sort(v[iter].begin(), v[iter].end(), greater<long long>());
    }
    priority_queue<Data> que;
    set<Data> se;
    que.push(Data(v[0][0] + v[1][0] + v[2][0], vector<int>({0, 0, 0})));
    for (int k = 0; k < K; ++k) {
        auto c = que.top();que.pop();
        cout << c.first << endl;
        
        // 次の三候補
        for (int iter = 0; iter < 3; ++iter) {
            if (c.second[iter] + 1 < N[iter]) {
                long long sum = c.first - v[iter][c.second[iter]] + v[iter][c.second[iter] + 1];
                auto num = c.second; num[iter]++;
                auto d = Data(sum, num);
                
                // すでに push されたもの以外
                if (!se.count(d)) se.insert(d), que.push(d);
            }
        }
    }
}

解法 2: 探索候補を絞る考え方 (editorial 解法 1)

全部の組合せは $XYZ$ 通りあって、これをすべて試していては間に合わないのだが、実はすごく上手い方法がある。

まず、 $A, B$ の各要素の合計値としてありうるものは $XY$ 個列挙することができる。そしてこの $XY$ 個の値のうち、上位 $K$ 個にも入らないやつは、合計した後に上位 $K$ 個以内に入ることはありえない。よって上位 $K$ 個のみを残して残りは捨ててしまってよい！！！

よって、さらに

$A, B$ の各要素の合計値としてありうるものの上位 $K$ 個
$C$ の各要素

との全組合せを改めて列挙することができて、それらの上位 $K$ 個を求めれば良い。計算量は $O(K^{2})$ 個の要素をソートする部分がボトルネックになっていて、 $O(K^{2}\log{K})$ となる。

ちなみに、このような感じで探索範囲を絞る考え方をする類題として、以下の問題がある。

atcoder.jp

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int main() {
    int X, Y, Z, K; cin >> X >> Y >> Z >> K;
    vector<long long> A(X), B(Y), C(Z);
    for (int i = 0; i < X; ++i) cin >> A[i];
    for (int i = 0; i < Y; ++i) cin >> B[i];
    for (int i = 0; i < Z; ++i) cin >> C[i];
    
    // A + B の上位 K 個
    vector<long long> AB;
    for (int i = 0; i < X; ++i) for (int j = 0; j < Y; ++j) AB.push_back(A[i] + B[j]);
    sort(AB.begin(), AB.end(), greater<long long>());
    AB.resize(K);
    
    // AB と C
    vector<long long> ABC;
    for (int i = 0; i < AB.size(); ++i) for (int j = 0; j < Z; ++j) ABC.push_back(AB[i] + C[j]);
    sort(ABC.begin(), ABC.end(), greater<long long>());
    
    // 出力
    for (int k = 0; k < K; ++k) cout << ABC[k] << endl;
}

さらに

もっと色々解法があって、すごく面白い！！！！！！
editorial がメッチャ丁寧！！！！！！

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 123 D - Cake 123 (水色, 400 点)

問題概要

制約

考えたこと (editorial 解法 3)

解法 2: 探索候補を絞る考え方 (editorial 解法 1)

さらに