平均がターゲット以上となる長さKのサブリストの数を求めるPythonプログラム

問題の概要

リスト nums と、2つの値 k および target が与えられたとします。このとき、「要素数がちょうど k であり、その平均値が target 以上であるサブリスト」の個数を求めるのが今回の課題です。

たとえば、入力が nums = [1, 10, 5, 6, 7]、k = 3、target = 6 の場合を考えてみましょう。このとき出力は 2 となります。サブリスト [1, 10, 7] の平均値は 6、サブリスト [10, 5, 6] の平均値は 7 となり、条件を満たすものが2つ存在するためです。

アルゴリズムの考え方

この問題は「スライディングウィンドウ（移動窓）」という手法を使うことで効率的に解くことができます。手順は以下の通りです。

• 比較のたびに除算を行わないよう、あらかじめ target *= k として「合計値との比較」に置き換えます（平均 ≥ target ⟺ 合計 ≥ target × k）。
• 合計を保持する変数 sum と、答えを格納する ans をそれぞれ 0 で初期化します。
• リストの各要素 n（インデックス i）に対して以下を処理します。
  • i >= k のとき、ウィンドウから外れる要素 nums[i - k] を合計から減算します。
  • 現在の要素 n を合計に加算します。
  • i >= k - 1（ウィンドウが k 個分埋まった状態）のとき、sum >= target であれば ans を 1 増やします。
• 最後に ans を返します。

実装例

それでは、以下のコードを見ながら理解を深めていきましょう。

サンプルコード

class Solution:
    def solve(self, nums, k, target):
        target *= k
        sum = 0
        ans = 0
        for i, n in enumerate(nums):
            if i >= k:
                sum -= nums[i - k]
            sum += n
            if i >= (k - 1):
                if sum >= target:
                    ans += 1
        return ans

ob = Solution()
nums = [1, 10, 5, 6, 7]
k = 3
target = 6
print(ob.solve(nums, k, target))

入力

[1, 10, 5, 6, 7], 3, 6

出力

2

計算量について

このアルゴリズムでは各要素を一度ずつ処理するため、時間計算量は O(n)、追加のメモリ使用量は O(1) で済みます。すべてのサイズ k のサブリストを毎回ゼロから合計し直す素朴な方法（O(n × k)）と比べて大幅に効率的であり、大きなデータセットでも高速に動作する点がメリットです。