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

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

AtCoder ABC 123 D - Cake 123 (400 点)

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 がメッチャ丁寧!!!!!!