LeetCode-in-Java

3397. Maximum Number of Distinct Elements After Operations

Medium

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:

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:

Solution

import java.util.Arrays;

public class Solution {
    public int maxDistinctElements(int[] nums, int k) {
        Arrays.sort(nums);
        int next = nums[0] - k + 1;
        int n = nums.length;
        int ans = 1;
        for (int i = 1; i < n; i++) {
            if (nums[i] + k < next) {
                continue;
            }
            next = Math.max(next, nums[i] - k);
            ans++;
            next++;
        }
        return ans;
    }
}