Problem Statement in English

Given an array of integers nums and an integer k, you can perform the following operation on the array:

  • Choose an element x from the array and either add or subtract k from it.

Your goal is to maximize the number of distinct elements in the array after performing any number of operations.

Approach

We can use a greedy approach to solve this problem.

The idea is to sort the array and then try to assign each number to the smallest possible value within its allowed range [num - k, num + k] that is greater than the last assigned value.

This way, we can maximize the number of distinct elements.

Solution in Python


class Solution:
    def maxDistinctElements(self, nums: List[int], k: int) -> int:
        last = -k
        res = 0
        nums.sort()

        for num in nums:
            c = max(num - k, last + 1)
            if c <= num + k:
                last = c
                res += 1

        return res

Complexity

  • Time: $O(n \log n)$
    Since we sort the array of size n.

  • Space: $O(1)$
    We use only a constant amount of extra space.

Mistakes I Made

I overcomplicated the implementation and had some failed test cases.

And we are done.