C++で1からkまでのすべての数を作るために追加すべき数の個数を求めるプログラム

問題概要

数値のリスト nums と整数 k が与えられます。nums の部分集合を選んで合計することで、1 から k までの任意の整数をすべて作れるようにしたいと考えます。そのために nums に新たに挿入する必要がある数の最小個数を求めてください。

たとえば、入力が nums = [3, 5]、k = 6 の場合、出力は 2 になります。「1」と「2」を挿入すればよいからです。実際、各値は次のように作成できます。

• 1 = [1]
• 2 = [2]
• 3 = [3]
• 4 = [1, 3]
• 5 = [5]
• 6 = [1, 5]

解き方(貪欲法)

この問題は貪欲法で効率的に解けます。ポイントは、「既存の数を使って 1 から sum まですべて作れる」という状態を管理することです。次に処理しようとする数が sum + 1 より大きい場合、sum + 1 を作る手段が存在しないため、sum + 1 という値そのものを追加するしかありません。

アルゴリズムの手順は以下の通りです。

• 配列 nums を昇順にソートする
• sum := 0、next := 1、ret := 0 で初期化する(sum は作れる範囲、next は次に必要な値)
• nums の各要素 i について次を繰り返す
  • next < i の間、次を繰り返す:
    • sum >= k ならループを抜ける
    • sum := sum + next とする
    • next := sum + 1 とする
    • ret を 1 増やす
  • sum >= k ならループを抜ける
  • sum := sum + i、next := sum + 1 とする
• next <= k の間、次を繰り返す:
  • sum := sum + next とする
  • next := sum + 1 とする
  • ret を 1 増やす
• ret を返す

なぜこの方法が機能するのか

1 〜 sum まで作れる状態で新しい数 x(x <= sum + 1)を追加すると、作れる範囲は 1 〜 (sum + x) まで拡張します。逆に、次の要素が sum + 1 より大きい場合は sum + 1 だけがどうしても作れないため、そこで sum + 1 を挿入します。これにより到達範囲は 1 〜 (2×sum + 1) まで一気に広がります。常に最も小さい欠落を埋めるというこの戦略により、挿入回数を最小に保てます。

C++での実装例

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

class Solution {
public:
    int solve(vector<int>& nums, int k) {
        sort(nums.begin(), nums.end());
        int sum = 0;
        int next = 1;
        int ret = 0;
        for (int i : nums) {
            while (next < i) {
                if (sum >= k) break;
                sum += next;
                next = sum + 1;
                ret++;
            }
            if (sum >= k) break;
            sum += i;
            next = sum + 1;
        }
        while (next <= k) {
            sum += next;
            next = sum + 1;
            ret++;
        }
        return ret;
    }
};

int solve(vector<int>& nums, int k) {
    return (new Solution())->solve(nums, k);
}

int main() {
    vector<int> v = {3, 5};
    int k = 6;
    cout << solve(v, k);
}

入力

[3, 5], 6

出力

2

計算量

ソートに O(n log n) かかります。その後の処理は、要素の走査 O(n) と挿入処理(next がほぼ倍々に増加するため最大でも O(log k) 回)で構成されるため、全体の計算量は O(n log n) となります。