Maximum Number of Distinct Elements After Operations

Med

#3018Time: O((2k+1)^N), where N is the length of `nums`. The search space is enormous as each of the N elements has up to `2k+1` choices. This is not a feasible solution.Space: O(N), for the recursion stack depth and to store the set of used values in a single path.

Patterns

Greedy

Algorithms

Sorting

Data structures

Array

Prompt

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:

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⁵
1 <= nums[i] <= 10⁹
0 <= k <= 10⁹

Approaches

2 approaches with complexity analysis and trade-offs.

This approach attempts to solve the problem by exploring all possible ways to assign new values to the elements of nums. For each element, we can change it by adding any integer from -k to k. A recursive function is used to traverse this vast search space of choices to find the assignment that results in the maximum number of distinct elements.

Algorithm

Define a recursive function, say solve(index, used_values).
index is the current index in the nums array we are considering.
used_values is a set containing the distinct values assigned to elements from nums[0] to nums[index-1].
The base case for the recursion is when index == nums.length. At this point, we have made assignments for all elements, and the number of distinct values is used_values.size(). We return this size.
In the recursive step for nums[index], we explore all possible valid assignments. For each possible new value v in the range [nums[index] - k, nums[index] + k], we can either add it to the set if it's new or reuse it if it's already present.
The function would look something like this: function solve(index, used_values):
- If index == nums.length, return used_values.size().
- max_distinct = 0
- For d from -k to k:
  - new_val = nums[index] + d
  - max_distinct = max(max_distinct, solve(index + 1, used_values U {new_val}))
- Return max_distinct.
The initial call would be solve(0, new HashSet<>()).

Walkthrough

We define a recursive function, say solve(index, used_values), where index is the current element nums[index] we are considering, and used_values is a set containing the distinct values assigned to elements from nums[0] to nums[index-1]. The base case for the recursion is when index reaches the end of the array. At this point, we have made assignments for all elements, and the number of distinct values is used_values.size(). We keep track of the maximum size found across all recursive paths.

In the recursive step for nums[index], we would try to assign it a new value. The new value v must be in the range [nums[index] - k, nums[index] + k]. We would have to iterate through all possible values v in this range, add v to a copy of the used_values set, and recurse with solve(index + 1, new_used_values). The final answer is the maximum size returned by any of these recursive explorations.

This method is a direct translation of the problem statement into a search algorithm. However, the number of choices at each step is 2k + 1, leading to a total number of states of roughly (2k+1)^N, which is computationally intractable for the given constraints.

Complexity

Time

O((2k+1)^N), where N is the length of `nums`. The search space is enormous as each of the N elements has up to `2k+1` choices. This is not a feasible solution.

Space

O(N), for the recursion stack depth and to store the set of used values in a single path.

Trade-offs

Pros

It is a correct formulation of the problem that theoretically explores all possibilities to find the true maximum.

Cons

Extremely inefficient and infeasible for the given constraints.
The range of k can be very large, making it impossible to iterate through all possible modifications d from -k to k.

The problem can be solved efficiently using a greedy strategy. The core idea is to process the numbers in a specific order and make a locally optimal choice at each step that leads to a globally optimal solution. By sorting the numbers first, we can ensure that when we assign a new value to an element, we do so in a way that minimally constrains the choices for subsequent, larger elements. We greedily pick the smallest possible new value for each element to leave the maximum 'room' for the elements that come after it.

Algorithm

1. Sort the input array nums in non-decreasing order.
1. Initialize a counter for distinct elements, distinct_count = 0.
1. Initialize a variable last_assigned_val to a value that is smaller than any possible result (e.g., a very small negative number like Long.MIN_VALUE / 2 to avoid overflow issues with +1).
1. Iterate through each num in the sorted nums array:
- a. Determine the smallest possible integer value (target) we can assign to num that is greater than last_assigned_val. This value is target = Math.max((long)num - k, last_assigned_val + 1).
- b. Check if this target value is achievable. The value must be within the allowed range for num, which is [num - k, num + k]. So, we check if target <= (long)num + k.
- c. If the target is achievable:
  - i. We have successfully found a new unique value. Increment distinct_count.
  - ii. Update last_assigned_val = target to ensure the next value we pick is even larger.
1. After iterating through all numbers, return distinct_count.

Walkthrough

First, we sort the input array nums in non-decreasing order. This allows us to handle the elements from smallest to largest, which is crucial for the greedy strategy.

We then iterate through the sorted array. We maintain a variable, last_assigned_val, which keeps track of the largest value we have assigned so far. For each number num in the sorted array, we try to assign it the smallest possible integer value that is (1) greater than last_assigned_val and (2) achievable from num (i.e., within the range [num - k, num + k]).

The smallest integer greater than last_assigned_val is last_assigned_val + 1. The smallest value we can form from num is num - k. Thus, the smallest valid candidate for our new value is target = Math.max((long)num - k, last_assigned_val + 1).

Next, we must check if this target value is possible. It is possible if it's not greater than the maximum value we can form from num, which is num + k. So, we check if target <= (long)num + k.

If this condition is true, we have found a new distinct value. We increment our count of distinct elements and update last_assigned_val to target. This ensures that the next element we process will be assigned an even larger value.

If the condition is false, it means there is no integer we can assign to num that is both greater than last_assigned_val and within its allowed range. In this case, we cannot make the current element distinct in this greedy manner, so we move on.

After the loop finishes, the total count represents the maximum number of distinct elements.

Complexity

Time

O(N log N), where N is the number of elements in `nums`. The sorting step dominates the complexity. The subsequent greedy pass through the array takes O(N) time.

Space

O(log N) or O(N), depending on the space complexity of the sorting algorithm used. For instance, in-place quicksort uses O(log N) space on average for the recursion stack, while mergesort uses O(N) auxiliary space.

Trade-offs

Pros

Highly efficient and provides the optimal solution.
The logic is straightforward and relatively easy to implement correctly.

Cons

The O(N log N) time complexity is determined by the sorting step, which is the most expensive part of the algorithm.

Solutions

Solution

1class Solution {2public3  int maxDistinctElements(int[] nums, int k) {4    Arrays.sort(nums);5    int n = nums.length;6    int ans = 0, pre = Integer.MIN_VALUE;7    for (int x : nums) {8      int cur = Math.min(x + k, Math.max(x - k, pre + 1));9      if (cur > pre) {10        ++ans;11        pre = cur;12      }13    }14    return ans;15  }16}

1class Solution {2public:3  int maxDistinctElements(vector<int> &nums, int k) {4    ranges ::sort(nums);5    int ans = 0, pre = INT_MIN;6    for (int x : nums) {7      int cur = min(x + k, max(x - k, pre + 1));8      if (cur > pre) {9        ++ans;10        pre = cur;11      }12    }13    return ans;14  }15};

1class Solution:2    def maxDistinctElements(self, nums: List[int], k: int) -> int: nums . sort() ans = 0 pre = - inf for x in nums: cur = min(x + k, max(x - k, pre + 1)) if cur > pre: ans += 1 pre = cur return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem