Count of triplets that can be removed without changing Mean of given Array

Last Updated : 12 Jan, 2023

Given an array arr[], the task is to calculate the count of possible triplets such that they can be removed from the array without changing the arithmetic mean of the array.

Example:

Input: arr[] = {8, 7, 4, 6, 3, 0, 7}
Output: 3
Explanation: The given array has 3 possible triplets such that removing them will not affect the arithmetic mean of the array. There are {7, 3, 0}, {4, 6, 0} and {3, 0, 7}.

Input: arr[] = {5, 5, 5, 5}
Output: 4

Approach: The given problem can be solved using the observation that for the mean of the remaining array to be constant, the mean of the removed triplet must be equal to the mean of the initial array. Hence the given problem is reduced to finding the count of triplets with the given sum which can be solved using hashing by following the below steps:

Iterate the given array arr[] for all possible values of pairs (a, b) and insert their sum into a map.
While iterating the array, check if (TargetSum – (a + b)) already exists in the map. If yes, then increment the value of the required count by its frequency.

Below is the implementation of the above approach:

C++

// C++ code for the above approach
#include <bits/stdc++.h>
using namespace std;
 
// Function to count the number of
// triplets with the given sum
int countTriplets(int arr[], int n, int sum)
{
    // Stores the final count
    int cnt = 0;
 
    // Map to store occurred elements
    unordered_map<int, int> m;
 
    for (int i = 0; i < n - 1; i++) {
        for (int j = i + 1; j < n; j++) {
 
            // Check if Sum - (a + b)
            // is present in map
            int k = sum - (arr[i] + arr[j]);
            if (m.find(k) != m.end())
 
                // Increment count
                cnt += m[k];
        }
 
        // Store the occurrences
        m[arr[i]]++;
    }
 
    // Return Answer
    return cnt;
}
 
// Function to C=find count of triplets
// that can be removed without changing
// arithmetic mean of the given array
int count_triplets(int arr[], int n)
{
    // Stores sum of all elements
    // of the given array
    int sum = 0;
 
    // Calculate the sum of the array
    for (int i = 0; i < n; i++) {
        sum = sum + arr[i];
    }
    // Store the arithmetic mean
    int mean = sum / n;
    int reqSum = 3 * mean;
 
    if ((3 * sum) % n != 0)
        return 0;
 
    // Return count
    return countTriplets(arr, n, reqSum);
}
 
// Driver Code
int main()
{
    int arr[] = { 5, 5, 5, 5 };
    int N = sizeof(arr) / sizeof(arr[0]);
    cout << count_triplets(arr, N);
 
    return 0;
}

Java

// Java code for the above approach
import java.util.HashMap;
 
class GFG {
 
  // Function to count the number of
  // triplets with the given sum
  static int countTriplets(int[] arr, int n, int sum)
  {
    // Stores the final count
    int cnt = 0;
 
    // Map to store occurred elements
    HashMap<Integer, Integer> m = new HashMap<>();
 
    for (int i = 0; i < n - 1; i++) {
      for (int j = i + 1; j < n; j++) {
 
        // Check if Sum - (a + b)
        // is present in map
        int k = sum - (arr[i] + arr[j]);
        if (m.containsKey(k))
 
          // Increment count
          cnt += m.get(k);
      }
 
      // Store the occurrences
      if (m.containsKey(arr[i]))
        m.put(arr[i],m.get(arr[i])+1);
      else
        m.put(arr[i],1);
    }
 
    // Return Answer
    return cnt;
  }
 
  // Function to C=find count of triplets
  // that can be removed without changing
  // arithmetic mean of the given array
  static int count_triplets(int[] arr, int n)
  {
    // Stores sum of all elements
    // of the given array
    int sum = 0;
 
    // Calculate the sum of the array
    for (int i = 0; i < n; i++) {
      sum = sum + arr[i];
    }
    // Store the arithmetic mean
    int mean = sum / n;
    int reqSum = 3 * mean;
 
    if ((3 * sum) % n != 0)
      return 0;
 
    // Return count
    return countTriplets(arr, n, reqSum);
  }
 
  // Driver Code
  public static void main (String[] args)
  {
    int[] arr = { 5, 5, 5, 5 };
    int N = arr.length;
    System.out.println(count_triplets(arr, N));
  }
}
 
// This code is contributed by Shubham Singh.

Python3

# python code for the above approach
 
# Function to count the number of
# triplets with the given sum
def countTriplets(arr, n, sum):
 
    # Stores the final count
    cnt = 0
 
    # Map to store occurred elements
    m = {}
 
    for i in range(0, n-1):
        for j in range(i+1, n):
 
            # Check if Sum - (a + b)
            # is present in map
            k = sum - (arr[i] + arr[j])
           
            if (k in m):
 
                # Increment count
                cnt += m[k]
 
        # Store the occurrences
        if arr[i] in m:
            m[arr[i]] += 1
        else:
            m[arr[i]] = 1
 
    # Return Answer
    return cnt
 
# Function to C=find count of triplets
# that can be removed without changing
# arithmetic mean of the given array
def count_triplets(arr, n):
 
    # Stores sum of all elements
    # of the given array
    sum = 0
 
    # Calculate the sum of the array
    for i in range(0, n):
        sum = sum + arr[i]
 
    # Store the arithmetic mean
    mean = sum // n
    reqSum = 3 * mean
 
    if ((3 * sum) % n != 0):
        return 0
 
    # Return count
    return countTriplets(arr, n, reqSum)
 
# Driver Code
if __name__ == "__main__":
 
    arr = [5, 5, 5, 5]
    N = len(arr)
    print(count_triplets(arr, N))
 
    # This code is contributed by rakeshsahni

C#

// C# code for the above approach
using System;
using System.Collections.Generic;
class GFG {
 
    // Function to count the number of
    // triplets with the given sum
    static int countTriplets(int[] arr, int n, int sum)
    {
        // Stores the final count
        int cnt = 0;
 
        // Map to store occurred elements
        Dictionary<int, int> m = new Dictionary<int, int>();
 
        for (int i = 0; i < n - 1; i++) {
            for (int j = i + 1; j < n; j++) {
 
                // Check if Sum - (a + b)
                // is present in map
                int k = sum - (arr[i] + arr[j]);
                if (m.ContainsKey(k))
 
                    // Increment count
                    cnt += m[k];
            }
 
            // Store the occurrences
            if (m.ContainsKey(arr[i]))
                m[arr[i]]++;
            else
                m[arr[i]] = 1;
        }
 
        // Return Answer
        return cnt;
    }
 
    // Function to C=find count of triplets
    // that can be removed without changing
    // arithmetic mean of the given array
    static int count_triplets(int[] arr, int n)
    {
        // Stores sum of all elements
        // of the given array
        int sum = 0;
 
        // Calculate the sum of the array
        for (int i = 0; i < n; i++) {
            sum = sum + arr[i];
        }
        // Store the arithmetic mean
        int mean = sum / n;
        int reqSum = 3 * mean;
 
        if ((3 * sum) % n != 0)
            return 0;
 
        // Return count
        return countTriplets(arr, n, reqSum);
    }
 
    // Driver Code
    public static void Main()
    {
        int[] arr = { 5, 5, 5, 5 };
        int N = arr.Length;
        Console.WriteLine(count_triplets(arr, N));
    }
}
 
// This code is contributed by ukasp.

Javascript

<script>
    // JavaScript code for the above approach
 
    // Function to count the number of
    // triplets with the given sum
    const countTriplets = (arr, n, sum) => {
     
        // Stores the final count
        let cnt = 0;
 
        // Map to store occurred elements
        let m = {};
 
        for (let i = 0; i < n - 1; i++) {
            for (let j = i + 1; j < n; j++) {
 
                // Check if Sum - (a + b)
                // is present in map
                let k = sum - (arr[i] + arr[j]);
                if (k in m)
 
                    // Increment count
                    cnt += m[k];
            }
 
            // Store the occurrences
            if (arr[i] in m) m[arr[i]]++;
            else m[arr[i]] = 1;
        }
 
        // Return Answer
        return cnt;
    }
 
    // Function to C=find count of triplets
    // that can be removed without changing
    // arithmetic mean of the given array
    const count_triplets = (arr, n) => {
     
        // Stores sum of all elements
        // of the given array
        let sum = 0;
 
        // Calculate the sum of the array
        for (let i = 0; i < n; i++) {
            sum = sum + arr[i];
        }
         
        // Store the arithmetic mean
        let mean = parseInt(sum / n);
        let reqSum = 3 * mean;
 
        if ((3 * sum) % n != 0)
            return 0;
 
        // Return count
        return countTriplets(arr, n, reqSum);
    }
 
    // Driver Code
    let arr = [5, 5, 5, 5];
    let N = arr.length;
    document.write(count_triplets(arr, N));
 
// This code is contributed by rakeshsahni
 
</script>

Output:

Time Complexity: O(N²), since there runs a nested loop both for N times.
Auxiliary Space: O(N), since N extra space is taken for hashing values.