すべての要素がK以上になるまで配列の要素を追加するC++プログラム｜最小ヒープによる効率的な解法

ソートされていない整数の配列 arr[] と整数 K が与えられたとき、配列内の2つの要素を選んで足し合わせて1つの要素にする操作を繰り返し、すべての要素を K 以上にするまでに必要な最小の操作回数を求めるのが本記事のテーマです。

問題の例

Input: arr[] = {1 10 12 9 2 3}, K = 6
Output: 2

解説

まず (1 + 2) を加算すると、新しい配列は 3 10 12 9 3 になります。次に (3 + 3) を加算すると、新しい配列は 6 10 12 9 となります。この時点で、リスト内のすべての要素が 6 以上になっていることが確認できます。

したがって、答えは 2、つまり 2 回の操作で条件を満たせることになります。

アプローチ：最小ヒープ（Min Heap）の活用

毎回配列全体から最小値を線形探索すると非効率なため、最小ヒープを使用します。最小ヒープを使えば、最小値の参照・挿入・削除をそれぞれ O(log n) で高速に行えます。

• 与えられた配列から最小ヒープを構築します。
• ヒープの最小値が K 未満である間、次の手順を繰り返します。
  • 最小の2つの要素を取り出し、その合計値を新しい要素としてヒープに挿入する。
  • 操作回数を1増やす。
• 要素が1つだけ残っていても K 未満の場合は -1 を返します（どのように操作しても条件を満たせないことを意味します）。

C++での実装コード

#include <bits/stdc++.h>
using namespace std;
class MinHeap {
    int *harr;
    int capacity;
    int heap_size;
    public:
    MinHeap(int *arr, int capacity);
    void heapify(int);
    int parent(int i) {
        return (i-1)/2;
    }
    int left(int i) {
        return (2*i + 1);
    }
    int right(int i) {
        return (2*i + 2);
    }
    int extractMin();
    int getMin() {
        return harr[0];
    }
    int getSize() {
        return heap_size;
    }
    void insertKey(int k);
};
MinHeap::MinHeap(int arr[], int n) {
    heap_size = n;
    capacity = n;
    harr = new int[n];
    for (int i=0; i<n; i++)
        harr[i] = arr[i];
    for (int i=n/2-1; i>=0; i--)
        heapify(i);
}
void MinHeap::insertKey(int k) {
    heap_size++;
    int i = heap_size - 1;
    harr[i] = k;
    while (i != 0 && harr[parent(i)] > harr[i]) {
        swap(harr[i], harr[parent(i)]);
        i = parent(i);
    }
}
int MinHeap::extractMin() {
    if (heap_size <= 0)
        return INT_MAX;
    if (heap_size == 1) {
        heap_size--;
        return harr[0];
    }
    int root = harr[0];
    harr[0] = harr[heap_size-1];
    heap_size--;
    heapify(0);
    return root;
}
void MinHeap::heapify(int i) {
    int l = left(i);
    int r = right(i);
    int smallest = i;
    if (l < heap_size && harr[l] < harr[i])
        smallest = l;
    if (r < heap_size && harr[r] < harr[smallest])
        smallest = r;
    if (smallest != i) {
        swap(harr[i], harr[smallest]);
        heapify(smallest);
    }
}
int countMinOps(int arr[], int n, int k) {
    MinHeap h(arr, n);
    long int res = 0;
    while (h.getMin() < k) {
        if (h.getSize() == 1)
            return -1;
        int first = h.extractMin();
        int second = h.extractMin();
        h.insertKey(first + second);
        res++;
    }
    return res;
}
int main() {
    int arr[] = {1, 10, 12, 9, 2, 3};
    int n = sizeof(arr)/sizeof(arr[0]);
    int k = 6;
    cout << countMinOps(arr, n, k);
    return 0;
}

出力結果

2

計算量について

ヒープの構築には O(n)、1回の操作ごとに要素の取り出しと挿入で O(log n) がかかります。操作は最大で n-1 回発生しうるため、全体の時間計算量は O(n log n) となります。これは素朴な線形探索による手法（O(n²)）と比べて大幅に効率的です。