【Python】1つの要素を削除して平均値をkにできるか判定する方法

問題概要

数値のリスト nums と整数 k が与えられたとき、リストからちょうど1つの要素を削除することで、残りの要素の平均値がちょうど k と一致するかどうかを判定します。この問題には、次の制約があります。

• 2 ≤ n ≤ 1,000(n は nums の要素数)
• nums[i] ≤ 1,000,000、k ≤ 1,000,000

たとえば、入力が [5,3,2,4,6,10]、k = 4 の場合、出力は True になります。10 を削除すると、残りの要素の平均は (5+3+2+4+6)/5 = 4 となり、k と一致するためです。

解法の考え方

この問題は、数学的な性質を利用することで効率的に解けます。要素 i を削除した後の平均が k になる条件を式で表すと、次のようになります。

(s − i) / (n − 1) = k

ここで、s は全要素の合計、n は要素数です。この式を変形すると s − i = k × (n − 1) となるため、各要素についてこの条件を満たすものが存在するかを確認すればよいことになります。具体的な手順は以下の通りです。

• s: nums の全要素の合計を求める
• t: k × (nums のサイズ − 1) を計算する
• nums の各要素 i について、s − i が t と等しければ True を返す
• すべての要素を確認しても該当がなければ False を返す

この手法の計算量は O(n) であり、リストを一度走査するだけで答えが得られるため、非常に効率的です。

実装例

理解を深めるために、以下の Python コードを見てみましょう。

class Solution:
    def solve(self, nums, k):
        s = sum(nums)
        t = k * (len(nums) - 1)
        for i in nums:
            if s - i == t:
                return True
        return False

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

入力

[5,3,2,4,6,10], 4

出力

True