Pythonでリストを1つの整数に減らすときの最小コストを求める方法

問題の概要

数値のリスト nums があるとします。このリストから任意の2つの数を選んで取り除き、その合計をリストの末尾に追加する操作を繰り返すことで、リストの長さを減らすことができます。この操作にかかるコストは、取り除いた2つの整数の合計です。ここで、nums を最終的に1つの整数にまで減らすときの合計コストの最小値を求めます。

入力例と動作の確認

例えば、nums = [2, 3, 4, 5, 6] の場合、出力は 45 になります。手順を順番に見てみましょう。

• まず 2 と 3 を取り除き、合計 5 を追加 → [4, 5, 6, 5](コスト 5)
• 次に 4 と 5 を取り除き、合計 9 を追加 → [6, 5, 9](コスト 9)
• 次に 6 と 5 を取り除き、合計 11 を追加 → [9, 11](コスト 11)
• 最後に 9 と 11 を取り除いて 19 になります(コスト 20)

コストの合計は 5 + 9 + 11 + 20 = 45 となります。

解法のアプローチ

この問題は「ハフマン符号化」と同じ構造を持つことで知られています。毎回できるだけ小さい2つの数を組み合わせることで、大きな数が何度も足し合わされるのを避けられます。そのため、常に最小の要素を優先的に取り出せるミニヒープ(優先度付きキュー)を使うのが最適です。

具体的な手順は以下の通りです。

• nums の要素でミニヒープを作成する
• ans := 0 と初期化する
• nums のサイズが 2 以上である限り、以下を繰り返す
  • a := ヒープから最小の要素を取り出す
  • b := 次に小さい要素を取り出す
  • ans := ans + a + b
  • a + b をヒープに挿入する
• ans を返す

Pythonでの実装例

class Solution:
    def solve(self, nums):
        import heapq
        heapq.heapify(nums)
        ans = 0
        while len(nums) >= 2:
            a = heapq.heappop(nums)
            b = heapq.heappop(nums)
            ans += a + b
            heapq.heappush(nums, a + b)
        return ans

ob = Solution()
nums = [2, 3, 4, 5, 6]
print(ob.solve(nums))

入力

[2, 3, 4, 5, 6]

出力

45

計算量について

heapify による初期化は O(n) で完了し、その後のループでは要素を1つ減らすごとに heappop を2回と heappush を1回行うため、全体の時間計算量は O(n log n)、空間計算量は O(n) となります。標準ライブラリの heapq を使うだけで効率的に実装できる点も、このアプローチの魅力です。