Pythonでリストの全要素を等しくするための最小総コストを求めるプログラム

nums と costs という2つの数値リストがあると仮定しましょう。ここで、nums[i] の値を costs[i] のコストで増加または減少させるという操作を考えます。この操作は何度でも実行でき、nums のすべての要素を同じ値に揃えたいとします。このとき、必要となる最小の総コストを求めるのが課題です。

たとえば、入力が nums = [3, 2, 4]、costs = [1, 10, 2] の場合、出力は 5 になります。これは、3 を 2 に減らすのにコスト 1 がかかり、さらに 4 を 2 回減らすのにそれぞれコスト 2 ずつ(合計 4)かかるためです。

解決のアプローチ

この問題を解くには、次の手順に従います。

• 引数として target を受け取る関数 helper() を定義します。

• total := 0 と初期化します。

• enumerate(nums) の各要素 i と n について、次の処理を行います。

  • target が n と異なる場合は、total := total + |n − target| × costs[i] を計算して加算します。

• total を返します。

• メイン処理では以下を実行します。

• low := 0、high := nums の最大値 とします。

• low < high である限り、次の処理を繰り返します。

  • mid := (low + high) / 2 とします。

  • helper(mid) < helper(mid + 1) であれば、high := mid とします。

  • それ以外の場合は、low := mid + 1 とします。

• 最後に helper(low) を返します。

なぜ二分探索が使えるのかというと、ある目標値 t に全要素を揃える際の総コストを表す関数 helper(t) は、t に関して下に凸(凸関数)の形になるためです。凸関数の最小値は二分探索によって効率よく絞り込むことができ、これにより全候補を総当たりするよりも高速に最適な目標値を見つけられます。

理解を深めるために、次の実装例を見てみましょう。

サンプルコード(Python)

class Solution:
    def solve(self, nums, costs):
        def helper(target):
            total = 0
            for i, n in enumerate(nums):
                if target != n:
                    total += abs(n - target) * costs[i]
            return total
        low, high = 0, max(nums)
        while low < high:
            mid = low + high >> 1
            if helper(mid) < helper(mid + 1):
                high = mid
            else:
                low = mid + 1
        return helper(low)

ob = Solution()
nums = [3, 2, 4]
costs = [1, 10, 2]
print(ob.solve(nums, costs))

入力

[3, 2, 4], [1, 10, 2]

出力

5