C++で2つの配列から最大数を作成するアルゴリズムを解説

問題概要

0〜9の数字からなり、長さ m と n をもつ2つの配列を考えます。それぞれの配列は1つの数値を表しています。この2つの配列から数字を選び出し、全体で k 桁となる最大の数を作成することを目指します。ただし重要な制約として、同じ配列から選んだ数字同士の相対的な順序は保持しなければなりません。

たとえば、入力が [3,4,7,5] と [9,1,3,5,8,4]、k = 5 の場合、求める答えは [9,8,7,5,4] となります。

解法のアプローチ

この問題は、次の3つの補助関数に分割すると見通しよく実装できます。

• modify(v, k)：1つの配列から、順序を保ちながら k 桁の最大部分列を抜き出します（単調スタックを使用）。
• mergeThem(nums1, nums2)：2つの部分列をマージし、最も大きくなる数列を構築します。
• greater(a, b, i, j)：2つの数列を先頭から辞書順に比較し、a の方が大きいかどうかを判定します。

さらにメイン処理では、「nums1 から i 桁、nums2 から k−i 桁」というすべての割り振りパターンについて候補を生成し、その中で最大のものを採用します。

手順1：mergeThem() — 2つの数列のマージ

• 配列 nums1 と nums2 を受け取ります。
• 結果格納用の配列 ret を用意します。
• i := 0、j := 0、n := nums1 のサイズ、m := nums2 のサイズ とします。
• i < n または j < m の間、次を繰り返します。
  • greater(nums1, nums2, i, j) が true の場合：ret の末尾に nums1[i] を追加し、i を1増やします。
  • それ以外の場合：ret の末尾に nums2[j] を追加し、j を1増やします。
• ret を返します。

手順2：modify() — 最大部分列の抽出

• 配列 v と桁数 k を受け取ります。
• スタック st と結果配列 ret を定義します。
• i := 0 から v のサイズ未満まで繰り返します。
  • x := v[i] とします。
  • 「スタックが空でない」「スタックの先頭要素が x より小さい」「st のサイズ ＋ v のサイズ − i − 1 ≥ k」という条件を満たす間、スタックから要素を取り除きます。
  • スタックのサイズが k 未満であれば、x をスタックに積みます。
• スタックが空になるまで、先頭要素を ret の末尾へ追加しながら取り除きます。
• ret を反転して返します。

手順3：greater() — 数列の大小比較

• 配列 a、b と開始インデックス i、j を受け取ります。
• i < a のサイズ かつ j < b のサイズ かつ a[i] == b[j] の間、i と j を1ずつ増やします。
• 「j == b のサイズ」または「i < a のサイズ かつ a[i] > b[j]」を満たす場合に true を返します。

手順4：メイン処理

• 結果配列 ret を定義します。
• n := nums1 のサイズ、m := nums2 のサイズ とします。
• i := 0 から k まで繰り返します。
  • i ≤ n かつ (k − i) ≤ m の場合、candidate = mergeThem(modify(nums1, i), modify(nums2, k − i)) として候補を生成します。
  • greater(candidate, ret, 0, 0) が true であれば、ret := candidate と更新します。
• ret を返します。

ポイント解説

modify() のスタック処理は貪欲法の一種です。「残りの要素数が十分ある限り、後から現れるより大きな数字のために小さな数字を捨てられる」という発想で、O(n) で最大部分列を求められます。また、マージ時に単純に先頭の大小だけを見ると誤りが生じるため、greater() で「以降の並び」まで辿って比較している点が工夫です。全体の時間計算量は、各 i の試行あたり O(m + n + k²)、試行回数が最大 k＋1 回であるため、おおよそ O(k(m + n + k²)) となります。

C++実装例

それでは、理解を深めるために実際のコードを見てみましょう。

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
        cout << v[i] << ", ";
    }
    cout << "]"<<endl;
}
class Solution {
public:
    vector<int> mergeThem(vector<int> nums1, vector<int> nums2)
    {
        vector<int> ret;
        int i = 0;
        int j = 0;
        int n = nums1.size();
        int m = nums2.size();
        while (i < n || j < m) {
            if (greater(nums1, nums2, i, j)) {
                ret.push_back(nums1[i]);
                i++;
            }
            else {
                ret.push_back(nums2[j]);
                j++;
            }
        }
        return ret;
    }
    vector<int> modify(vector<int>& v, int k)
    {
        stack<int> st;
        vector<int> ret;
        for (int i = 0; i < v.size(); i++) {
            int x = v[i];
            while (!st.empty() && st.top() < x && st.size() + (v.size() - i) - 1 >= k) {
                st.pop();
            }
            if (st.size() < k)
                st.push(x);
            }
            while (!st.empty()) {
                ret.push_back(st.top());
                st.pop();
            }
            reverse(ret.begin(), ret.end());
            return ret;
        }
        bool greater(vector<int>& a, vector<int>& b, int i, int j)
        {
            while (i < a.size() && j < b.size() && a[i] == b[j])
            i++, j++;
            return j == b.size() || (i < a.size() && a[i] > b[j]);
        }
        vector<int> maxNumber(vector<int>& nums1, vector<int>& nums2, int k)
        {
            vector<int> ret;
            int n = nums1.size();
            int m = nums2.size();
            for (int i = 0; i <= k; i++) {
                if (i <= n && (k - i) <= m) {
                    vector<int> candidate = mergeThem(modify(nums1, i), modify(nums2, k - i));
                    if (greater(candidate, ret, 0, 0)) {
                        ret = candidate;
                    }
                }
            }
            return ret;
        }
};
main() {
    Solution ob;
    vector<int> v = { 3, 4, 7, 5 }, v1 = { 9, 1, 3, 5, 8, 4 };
    print_vector(ob.maxNumber(v, v1, 5));
}

実行例

入力

{ 3, 4, 7, 5 }
{ 9, 1, 3, 5, 8, 4 }
5

出力

[9, 8, 7, 5, 4]