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.

Example 1

1
2
3
4
5

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

1
2
3
4
5

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

Examples

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

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();
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

import "sort"
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)
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

import java.util.*;
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();
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

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
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10

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)
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

use std::collections::HashSet;
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
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

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.