C++で複数の棒を1本に接続する際の最小コストを求めるアルゴリズム

問題の概要

正の整数の長さを持つ複数の棒があると仮定します。長さが X と Y の2本の棒を1本に接続するときのコストは X + Y です。この操作を、棒が1本だけ残るまで繰り返します。ここで、与えられたすべての棒をこの方法で1本にまとめるときの最小コストを求めることが目的です。

例えば、棒の長さの配列が [2, 4, 3] の場合、出力は 14 になります。実際の手順としては、まず 2 と 3 を接続して 5 を作り（コスト 5）、次に 5 と 4 を接続して 9 を作る（コスト 9）ので、合計コストは 5 + 9 = 14 となります。他の接続順序ではコストが大きくなるため、これが最適解です。

解決アプローチ：貪欲法と最小ヒープ

この問題は貪欲法で解くことができます。「常に現時点で最も短い2本の棒を選んで接続する」という戦略を繰り返すことで、全体のコストを最小化できます。短い棒ほど後の工程で何度も足し合わせられるため、小さい値から順に処理することが重要です。これはハフマン符号化と同じ発想に基づいています。

効率的に最小値を取り出すために、最小ヒープ（min-heap）として動作する優先度付きキューを使用します。手順は以下の通りです。

• 最小ヒープ型の優先度付きキュー pq を定義する
• s のすべての要素を pq に挿入する
• ans := 0 と初期化する
• pq の要素数が 2 以上である限り、以下を繰り返す：
  • pq の先頭（最小値）を取り出し、temp に代入する
  • pq の次の先頭（最小値）を temp に加算し、pq から取り出す
  • ans := ans + temp
  • temp を pq に挿入する
• ans を返す

なお、計算量はヒープへの挿入・削除がそれぞれ O(log n) であり、これを要素ごとに行うため、全体で O(n log n) となり非常に効率的です。

C++での実装例

以下の実装例を見ると、理解がより深まるでしょう。

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    int connectSticks(vector<int>& s) {
        priority_queue <int, vector<int>, greater<int> > pq;
        for(int i =0;i<s.size();i++)pq.push(s[i]);
        int ans = 0;
        while(pq.size()>1){
            int temp = pq.top();
            pq.pop();
            temp += pq.top();
            pq.pop();
            ans+=temp;
            pq.push(temp);
        }
        return ans;
    }
};
main(){
    vector<int> v = {2,4,3};
    Solution ob;
    cout <<ob.connectSticks(v);
}

ポイントは、priority_queue<int, vector<int>, greater<int> のように第3テンプレート引数に greater<int> を指定することです。これにより、通常の C++ の優先度付きキュー（最大ヒープ）とは逆の最小ヒープとして動作し、常に最も小さい値が先頭に来るようになります。

実行結果

入力：

[2,4,3]

出力：

14