Maximum Number of Distinct Elements After Operations

main.md

Question

Approach

The key idea to solve this problem efficiently is to treat each element x in the array as an interval [x - k, x + k], representing all the possible values it can take after performing the allowed operation. To maximize the number of distinct elements, we can use a greedy approach: sort all these intervals by their right endpoint, then iterate through them while maintaining a pointer cur that tracks the smallest integer not yet used. For each interval [l, r], we first move cur to the maximum of its current value and the interval’s left endpoint (cur = max(cur, l)), ensuring we stay within the valid range. If cur is still less than or equal to r, we can “use” this value — we increment our answer count and move cur forward by one (cur += 1) to avoid reusing numbers.

Implementation

class Solution:
    def maxDistinctElements(self, nums: List[int], k: int) -> int:
        intervals = [(x - k, x + k) for x in nums]
        intervals.sort(key=lambda it: (it[1], it[0]))

        next_avail = -10**30
        ans = 0
        for l, r in intervals:
            candidate = max(next_avail, l)
            if candidate <= r:
                ans += 1
                next_avail = candidate + 1
        return ans

Complexities

• Time: O(n logn)
• Space: O(n)