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

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

JOI 一次予選 2020 (第 1 回) C - マージ (5Q, 難易度 3)

マージソートのマージの部分を実装する問題ですね。
添字を順に進めていくような処理を実装します! この実装は、のちにしゃくとり法を学ぶ際の参考にもなります。

問題概要

制約

  •  1 \le N, M \le 500

解法

問題文が複雑で、何をすればいいのかを理解するのが大変かもしれないですね。

一般に、数列  A, B の先頭の要素を削除していく処理は実装が大変です (なお、数列を deque で実装するとか、数列を reverse して末尾から削除していくような実装をするなどの別解はあります)。

そこで、代わりに、数列  A 上の添字  i と、数列  B 上の添字  j を管理することにします。これらは「今この添字のある位置が疑似的に数列の先頭になっています」を表しています。

たとえば、上図のような状況を考えます。数列  A の左 2 個 (2, 4) と、数列  B の左 1 個 (3) がすでに削除されて、数列  C に格納されている状態とします。

次のステップでは、数列  A の真の先頭の要素 (疑似的に値 A[2] = 7) と、数列  B の先頭の要素 (疑似的に値 B[1] = 8) とを比較します。A[2] の方が小さいので、A[2] を数列  C に push して、 A の添字  i を 1 進めます。

このような処理を、添字  i, j がそれぞれ数列  A, B の終端に到達するまで繰り返せばよいでしょう。

コード

#include <bits/stdc++.h>
using namespace std;

int main() {
    int N, M;
    cin >> N >> M;
    vector<int> A(N), B(M);
    for (int i = 0; i < N; ++i) cin >> A[i];
    for (int i = 0; i < M; ++i) cin >> B[i];
    
    vector<int> C;
    int i = 0, j = 0;  // A 側の先頭の添字と B 側の先頭の添字
    while (i < N || j < M) {
        if (i == N) {
            // A 側が終端に達していたら B 側を入れて、B 側の添字 j を進める
            C.push_back(B[j]);
            ++j;
        } else if (j == M) {
            // B 側が終端に達していたら A 側を入れて、A 側の添字 i を進める
            C.push_back(A[i]);
            ++i;
        } else if (A[i] < B[j]) {
            // A の方が小さいとき、A 側を入れて、A 側の添字 i を進める
            C.push_back(A[i]);
            ++i;
        } else {
            // B の方が小さいとき、B 側を入れて、B 側の添字 j を進める
            C.push_back(B[j]);
            ++j;
        }
    }
    
    // 出力
    for (int i = 0; i < N + M; ++i) cout << C[i] << endl;
}