Remove One to Make Average k in Python

Suppose we have a list of numbers called nums and an integer k, we have to check whether we can remove exactly one element from the list to make the average equal to exactly k. Now we have to keep in mind that, there are some constraints ?

• 2 ? n ? 1,000 where n is number of elements of nums list
• nums[i], k ? 1,000,000

So, if the input is like [5,3,2,4,6,10], k = 4, then the output will be True as if we remove 10, then the average of elements will be (5+3+2+4+6)/5 = 4, which is same as k.

Algorithm

To solve this, we will follow these steps ?

• Calculate the total sum of all elements in nums
• For each element, check if removing it gives the desired average
• If (sum - element) / (length - 1) equals k, return True
• If no element satisfies the condition, return False

Mathematical Approach

For the average to be k after removing one element, we need:

(sum - element) / (length - 1) = k

Rearranging: sum - element = k × (length - 1)

Example

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

# Test the solution
ob = Solution()
nums = [5, 3, 2, 4, 6, 10]
k = 4
result = ob.solve(nums, k)
print(f"Can remove one element to make average {k}? {result}")

# Let's verify by checking which element to remove
s = sum(nums)
print(f"Original sum: {s}")
print(f"Target sum after removal: {k * (len(nums) - 1)}")

for i, num in enumerate(nums):
    if s - num == k * (len(nums) - 1):
        remaining = [x for j, x in enumerate(nums) if j != i]
        avg = sum(remaining) / len(remaining)
        print(f"Remove {num}: remaining = {remaining}, average = {avg}")

Can remove one element to make average 4? True
Original sum: 30
Target sum after removal: 20
Remove 10: remaining = [5, 3, 2, 4, 6], average = 4.0

Alternative Approach Without Class

def can_make_average_k(nums, k):
    """Check if we can remove exactly one element to make average equal to k"""
    total_sum = sum(nums)
    target_sum = k * (len(nums) - 1)
    
    return target_sum in [total_sum - num for num in nums]

# Test cases
test_cases = [
    ([5, 3, 2, 4, 6, 10], 4),
    ([1, 2, 3, 4, 5], 3),
    ([10, 20, 30], 15)
]

for nums, k in test_cases:
    result = can_make_average_k(nums, k)
    print(f"nums = {nums}, k = {k} ? {result}")

nums = [5, 3, 2, 4, 6, 10], k = 4 ? True
nums = [1, 2, 3, 4, 5], k = 3 ? True
nums = [10, 20, 30], k = 15 ? True

Time and Space Complexity

• Time Complexity: O(n) where n is the length of nums
• Space Complexity: O(1) for the class-based solution, O(n) for the alternative approach

Conclusion

The solution uses a mathematical approach to check if removing any single element results in the desired average. By comparing the target sum with the sum after removing each element, we can efficiently determine if it's possible to achieve the target average k.