Maximum Number of Distinct Elements After Operations

Problem

You are given an integer array nums and an integer k.

You are allowed to perform the following operation on each element of the array at most once :

• Add an integer in the range [-k, k] to the element.

Return the maximum possible number of distinct elements in nums after performing the operations.

Examples

Example 1

Input: nums = [1,2,2,3,3,4], k = 2
Output: 6
Explanation:
`nums` changes to `[-1, 0, 1, 2, 3, 4]` after performing operations on the
first four elements.

Example 2

Input: nums = [4,4,4,4], k = 1
Output: 3
Explanation:
By adding -1 to `nums[0]` and 1 to `nums[1]`, `nums` changes to `[3, 5, 4,
4]`.

Constraints

• 1 <= nums.length <= 10^5
• 1 <= nums[i] <= 10^9
• 0 <= k <= 10^9

Solution

Method 1 – Greedy + Set

Intuition

To maximize the number of distinct elements, for each number, try to move it as far as possible within the range [num-k, num+k] to avoid collisions with other numbers. We can greedily assign each number to the smallest available value in its range.

Approach

1. Sort the input array.
2. Use a set to keep track of used values.
3. For each number, try to assign it to the smallest value in [num-k, num+k] that is not already used.
4. Count the number of unique assignments.

Code

C++

class Solution {
public:
    int maxDistinctElements(vector<int>& nums, int k) {
        sort(nums.begin(), nums.end());
        unordered_set<int> used;
        for (int num : nums) {
            for (int x = num - k; x <= num + k; ++x) {
                if (!used.count(x)) {
                    used.insert(x);
                    break;
                }
            }
        }
        return used.size();
    }
};

Go

func maxDistinctElements(nums []int, k int) int {
    sort.Ints(nums)
    used := make(map[int]bool)
    for _, num := range nums {
        for x := num - k; x <= num + k; x++ {
            if !used[x] {
                used[x] = true
                break
            }
        }
    }
    return len(used)
}

Java

class Solution {
    public int maxDistinctElements(int[] nums, int k) {
        Arrays.sort(nums);
        Set<Integer> used = new HashSet<>();
        for (int num : nums) {
            for (int x = num - k; x <= num + k; x++) {
                if (!used.contains(x)) {
                    used.add(x);
                    break;
                }
            }
        }
        return used.size();
    }
}

Kotlin

class Solution {
    fun maxDistinctElements(nums: IntArray, k: Int): Int {
        nums.sort()
        val used = mutableSetOf<Int>()
        for (num in nums) {
            for (x in num - k..num + k) {
                if (x !in used) {
                    used.add(x)
                    break
                }
            }
        }
        return used.size
    }
}

Python

class Solution:
    def maxDistinctElements(self, nums: list[int], k: int) -> int:
        nums.sort()
        used = set()
        for num in nums:
            for x in range(num - k, num + k + 1):
                if x not in used:
                    used.add(x)
                    break
        return len(used)

Rust

impl Solution {
    pub fn max_distinct_elements(mut nums: Vec<i32>, k: i32) -> i32 {
        nums.sort();
        let mut used = HashSet::new();
        for &num in &nums {
            for x in num - k..=num + k {
                if !used.contains(&x) {
                    used.insert(x);
                    break;
                }
            }
        }
        used.len() as i32
    }
}

TypeScript

class Solution {
    maxDistinctElements(nums: number[], k: number): number {
        nums.sort((a, b) => a - b);
        const used = new Set<number>();
        for (const num of nums) {
            for (let x = num - k; x <= num + k; x++) {
                if (!used.has(x)) {
                    used.add(x);
                    break;
                }
            }
        }
        return used.size;
    }
}

Complexity

• ⏰ Time complexity: O(n * k), where n is the length of nums and for each number we may check up to 2k+1 values.
• 🧺 Space complexity: O(n), for the set of used values.

Method 2 - Greedy (Sort and assign the lowest possible distinct value)

Intuition

If we process values in increasing order and always assign the smallest distinct value available within an element's allowed interval [num - k, num + k], we minimize collisions with future elements. Sorting lets us fix the lowest elements first; tracking the last assigned distinct value ensures each new assignment is strictly greater than previous ones when possible.

Approach

1. Sort nums in non-decreasing order.
2. Maintain a variable lastAssigned that stores the last value we assigned to keep results distinct (initialize it to a very small sentinel).
3. For each num in nums compute the smallest feasible assignment low = max(num - k, lastAssigned + 1). This picks the lowest value in the allowed interval that is strictly greater than the previous assignment.
4. If low <= num + k, the value low is a valid assignment — set lastAssigned = low and increment the distinct count.
5. Return the count after processing all elements.

Complexity

• ⏰ Time complexity: O(n log n) – Sorting dominates the runtime; we then scan the array once.
• 🧺 Space complexity: O(1) – Only a few scalar variables are used (ignoring input storage).

Code

C++

class Solution {
public:
    int maxDistinctElements(std::vector<int>& nums, int k) {
        std::sort(nums.begin(), nums.end());
        long long lastAssigned = LLONG_MIN / 4;  // sentinel
        int count = 0;
        for (int x : nums) {
            long long low = std::max<long long>((long long)x - k, lastAssigned + 1);
            if (low <= (long long)x + k) {
                lastAssigned = low;
                ++count;
            }
        }
        return count;
    }
};

Go

func maxDistinctElements(nums []int, k int) int {
    sort.Ints(nums)
    const sentinel = int64(-1 << 60)
    lastAssigned := sentinel
    count := 0
    for _, num := range nums {
        low := maxInt64(int64(num)-int64(k), lastAssigned+1)
        if low <= int64(num)+int64(k) {
            lastAssigned = low
            count++
        }
    }
    return count
}

func maxInt64(a, b int64) int64 {
    if a > b {
        return a
    }
    return b
}

Java

class Solution {
    public int maxDistinctElements(int[] nums, int k) {
        Arrays.sort(nums);
        long lastAssigned = Long.MIN_VALUE / 4; // sentinel
        int count = 0;
        for (int num : nums) {
            long low = Math.max((long) num - k, lastAssigned + 1);
            if (low <= (long) num + k) {
                lastAssigned = low;
                count++;
            }
        }
        return count;
    }
}

Kotlin

class Solution {
    fun maxDistinctElements(nums: IntArray, k: Int): Int {
        nums.sort()
        var lastAssigned = Long.MIN_VALUE / 4
        var count = 0
        for (num in nums) {
            val low = maxOf(num.toLong() - k, lastAssigned + 1)
            if (low <= num.toLong() + k) {
                lastAssigned = low
                count++
            }
        }
        return count
    }
}

Python

class Solution:
    def maxDistinctElements(self, nums: list[int], k: int) -> int:
        nums.sort()
        last_assigned = -10**18
        count = 0
        for num in nums:
            low = max(num - k, last_assigned + 1)
            if low <= num + k:
                last_assigned = low
                count += 1
        return count

Rust

impl Solution {
    pub fn max_distinct_elements(mut nums: Vec<i32>, k: i32) -> i32 {
        nums.sort();
        let mut last_assigned: i64 = i64::MIN / 4; // sentinel
        let mut count: i32 = 0;
        for num in nums {
            let low = std::cmp::max(num as i64 - k as i64, last_assigned + 1);
            if low <= num as i64 + k as i64 {
                last_assigned = low;
                count += 1;
            }
        }
        count
    }
}

TypeScript

class Solution {
    maxDistinctElements(nums: number[], k: number): number {
        nums.sort((a, b) => a - b);
        let lastAssigned = Number.MIN_SAFE_INTEGER;
        let count = 0;
        for (const num of nums) {
            const low = Math.max(num - k, lastAssigned + 1);
            if (low <= num + k) {
                lastAssigned = low;
                count++;
            }
        }
        return count;
    }
}