Pythonで最大値と最小値の差が最小になるサイズkのリストを選ぶ方法

数値のリスト nums と整数 k が与えられたとき、nums から要素を選んでサイズ k のリストを作成し、そのリスト内の最大値と最小値の差をできるだけ小さくすることを考えます。そして、その最小の差を返すのがこの問題の目的です。

例えば、入力が nums = [3, 11, 6, 2, 9]、k = 3 の場合、出力は 4 になります。これは、最適なリストとして [2, 3, 6] を作成でき、最大値 6 と最小値 2 の差が 4 になるためです。

解法のアプローチ

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

• リスト nums をソートする

• 空のリスト ls を用意する

• i を 0 から len(nums) - k + 1 の範囲でループし、ls の末尾に nums[i + k - 1] - nums[i] を追加する

• ls の最小値を返す

このアプローチが有効なのは、ソート後のリストでは連続する k 個の要素を選ぶのが常に最適だからです。最大値と最小値の差を最小化するには、値が近い要素同士を選べばよく、ソート済みリスト上で隣接する k 個の組(スライディングウィンドウ)ごとに差を計算すれば、その中の最小値が答えになります。

実装例

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

class Solution:
    def solve(self, nums, k):
        nums.sort()
        ls = []
        for i in range(len(nums) - k + 1):
            ls.append(nums[i + k - 1] - nums[i])
        return min(ls)

ob = Solution()
nums = [3, 11, 6, 2, 9]
k = 3
print(ob.solve(nums, k))

入力

[3, 11, 6, 2, 9], 3

出力

4

計算量の目安

このアルゴリズムでは、ソートに O(n log n)、各ウィンドウの差の計算に O(n) のコストがかかるため、全体の時間計算量は O(n log n)(n はリストの要素数)となります。空間計算量は結果を格納するためのリスト分の O(n) です。シンプルでありながら効率的に動作する手法と言えます。