Counting Distinct elements after Range of Operations

Given an integer K and array A[] of size N, count the number of distinct elements after performing an operation on every array index. For each
i (0 ≤ i ≤ N – 1), we can perform one of the following operations exactly once:

• Add K to A[i]
• Subtract K from A[i]
• Leave A[i] as it is.

Examples:

Input: N = 6, K = 2, A[] = {2, 4, 3, 2, 2, 1}
Output: 6
Explanation: Subtract K from the first element, Add K to the second and fourth element and Leave the third, fifth and sixth element as it is. After performing these operations we will obtain A[] = {0, 6, 3, 4, 2, 1}. Hence, the number of distinct elements are 6.

Input: N = 5, K = 3, A[] = {3, 3, 3, 3, 3}
Output: 3
Explanation: Subtract K from first element, leave second element as it is and add K to the third element. For the remaining elements, it doesn’t matter which operation we perform they will remain duplicates.

Approach: This can be solved by the following approach:

The problem can be solved by using a Greedy Approach. Sort the array A[] in ascending order. Now, if we start from the smallest element we have to find what will be the most optimal operation to perform on any element. At any index i, we can perform three operations:

• Increase A[i] by K, but this can result to duplicate elements when we iterate over the next elements which are greater or equal to A[i]
• Leave A[i] as it is, but we might need to decrease the upcoming greater elements by K which can again lead to duplicate elements.
• Decrease A[i] by K, this is the best operation which we can perform if we have not encountered any element having value A[i] – K as decreasing A[i] will allow the upcoming greater elements to decrease themselves by K without any possibility of duplicate elements.

In conclusion, we should try to decrease A[i] by K, if possible, else leave A[i] as it is. If leaving A[i] as it is also leads to duplicate elements, then we can increase A[i] by K.

Steps to solve the above problem:

• Sort the array A[] in ascending order.
• Maintain a variable cnt to count the number of distinct elements and a map freq to store the frequency of elements
• Iterating in A array:
  • Check whether A[ i ] – K exists, if it doesn’t decrease A[i] by K and increment cnt by 1.
  • Else check whether A[ i ] exists, if it doesn’t then leave A[i] as it is and increment cnt by 1.
  • Check if A[i] + K exists, if it doesn’t increase A[i] by K and increment cnt by 1.
• Return the cnt.

Below is the implementation of the code:

C++

// C++ Implementation of the code

#include <bits/stdc++.h>
using namespace std;

// Function to find maximum distinct
// elements
int distinctElements(int N, int K, vector<int> A)
{

    // Sorting the array A
    sort(A.begin(), A.end());

    map<int, int> freq;
    int cnt = 0;

    for (int i = 0; i < N; i++) {

        // Checking A[j] - K exists
        if (freq[A[i] - K] == 0) {
            freq[A[i] - K] = 1;
            ++cnt;
        }

        // Checking A[j] exists
        else if (freq[A[i]] == 0) {
            freq[A[i]] = 1;
            ++cnt;
        }

        // Checking A[j] + K exists
        else if (freq[A[i] + K] == 0) {
            freq[A[i] + K] = 1;
            ++cnt;
        }
    }

    // Returning total distinct elements
    return cnt;
}

// Driver code
int main()
{
    int N = 6, K = 2;
    vector<int> A = { 2, 4, 3, 2, 2, 1 };

    // Function call
    cout << distinctElements(N, K, A) << endl;

    return 0;
}

Java

import java.util.*;

public class DistinctElements {
    // Function to find maximum distinct elements
    public static int distinctElements(int N, int K, List<Integer> A) {
        // Sorting the list A
        Collections.sort(A);

        Map<Integer, Integer> freq = new HashMap<>();
        int cnt = 0;

        for (int i = 0; i < N; i++) {
            // Checking A[j] - K exists
            if (freq.getOrDefault(A.get(i) - K, 0) == 0) {
                freq.put(A.get(i) - K, 1);
                ++cnt;
            }

            // Checking A[j] exists
            else if (freq.getOrDefault(A.get(i), 0) == 0) {
                freq.put(A.get(i), 1);
                ++cnt;
            }

            // Checking A[j] + K exists
            else if (freq.getOrDefault(A.get(i) + K, 0) == 0) {
                freq.put(A.get(i) + K, 1);
                ++cnt;
            }
        }

        // Returning total distinct elements
        return cnt;
    }

    // Driver code
    public static void main(String[] args) {
        int N = 6, K = 2;
        List<Integer> A = Arrays.asList(2, 4, 3, 2, 2, 1);

        // Function call
        System.out.println(distinctElements(N, K, A));
    }
}

Python3

def distinctElements(N, K, A):
    # Sorting the list A
    A.sort()

    freq = {}
    cnt = 0

    for i in range(N):
        # Checking A[j] - K exists
        if freq.get(A[i] - K, 0) == 0:
            freq[A[i] - K] = 1
            cnt += 1

        # Checking A[j] exists
        elif freq.get(A[i], 0) == 0:
            freq[A[i]] = 1
            cnt += 1

        # Checking A[j] + K exists
        elif freq.get(A[i] + K, 0) == 0:
            freq[A[i] + K] = 1
            cnt += 1

    # Returning total distinct elements
    return cnt

# Driver code
N = 6
K = 2
A = [2, 4, 3, 2, 2, 1]

# Function call
print(distinctElements(N, K, A))

# This code is contributed by akshitaguprzj3

C#

using System;
using System.Collections.Generic;
using System.Linq;

public class GFG
{
    // Function to find maximum distinct elements
    public static int DistinctElements(int N, int K, List<int> A)
    {
        // Sorting the list A
        A.Sort();

        Dictionary<int, int> freq = new Dictionary<int, int>();
        int cnt = 0;

        for (int i = 0; i < N; i++)
        {
            // Checking A[j] - K exists
            if (!freq.ContainsKey(A[i] - K) || freq[A[i] - K] == 0)
            {
                freq[A[i] - K] = 1;
                cnt++;
            }

            // Checking A[j] exists
            else if (!freq.ContainsKey(A[i]) || freq[A[i]] == 0)
            {
                freq[A[i]] = 1;
                cnt++;
            }

            // Checking A[j] + K exists
            else if (!freq.ContainsKey(A[i] + K) || freq[A[i] + K] == 0)
            {
                freq[A[i] + K] = 1;
                cnt++;
            }
        }

        // Returning total distinct elements
        return cnt;
    }

    // Driver code
    public static void Main(string[] args)
    {
        int N = 6, K = 2;
        List<int> A = new List<int> { 2, 4, 3, 2, 2, 1 };

        // Function call
        Console.WriteLine(DistinctElements(N, K, A));
    }
}

Javascript

// JavaScript Implementation of the code

// Function to find maximum distinct elements
function distinctElements(N, K, A) {
    // Sorting the array A
    A.sort((a, b) => a - b);

    // Create a Map to store the frequency
    // of elements with a specific difference (K)
    const freq = new Map();

    // Initialize a count for d
    // istinct elements
    let cnt = 0;

    // Iterate through the elements
    // in the array A
    for (let i = 0; i < N; i++) {

        // Checking if A[i] - K exists
        if (!freq.has(A[i] - K)) {
            freq.set(A[i] - K, 1);
            cnt++;
        }

        // Checking if A[i] exists
        else if (!freq.has(A[i])) {
            freq.set(A[i], 1);
            cnt++;
        }
        // Checking if A[i] + K exists
        else if (!freq.has(A[i] + K)) {
            freq.set(A[i] + K, 1);
            cnt++;
        }
    }

    // Return the total count
    // of distinct elements
    return cnt;
}

// Input values
const N = 6;
const K = 2;
const A = [2, 4, 3, 2, 2, 1];

// Function call and display the result
console.log(distinctElements(N, K, A));

Output

6

Time Complexity: O(N*log N), where N is the size of the array A[].
Auxiliary Space: O(N)