Count of pairs of Array elements with average at least K

Last Updated : 28 Mar, 2023

Given an array A[] of size N consisting of N integers, the task is to count the number of pairs such that their average is greater or equal to K.

Example:

Input: N = 4, K = 3, A = {5, 1, 3, 4}
Output: 4
Explanation: (5, 1), (5, 3), (5, 4) and (3, 4) are the required pairs with average greater or equal to K = 3.

Input: N = 3, K = 3, A = {1, 2, 3}
Output: 0
Explanation: No pairs exist with average greater or equal to K = 3.

Approach: This problem can be solved using binary search for the first occurrence of an element in based on the following observation:

Just need to find 2*K – A[i] because average of two numbers X and Y is K ≤ (X+Y)/2. Now replace X with the current element that we are traversing i.e A[i] then equation becomes Y ≥ 2*K-A[i].

So for any element A[i] in the array A[] total number of pairs formed are all the numbers in A[] that are greater than or equal to 2*K-A[i] i.e size of array ‘A’ -index of 2*K-A[i].

So go as left as possible in A[] and for that find the first occurrence of 2*K-A[i]. If 2*K-A[i] is not found in A[] then return the index of next greater element of 2*K-A[i] because if average ≤ (X+Y)/2 for any two integers then also average ≤ (X+Z)/2 for all Z ≥ Y.

Follow the below steps to solve this problem:

Sort the array A[].
Traverse the array A[].
- Find the first occurrence of 2*k-A[i] for every element A[i] in A.
- If 2*k-A[i] does not exist then find the first occurrence of an element just greater than 2*k-A[i] in the array A and store its result in a variable (say ind).
- If ind is not -1, then add N-ind in the answer.
Return the final answer after the traversal is over.

Below is the implementation of the above approach:

C++14

// C++ code to implement the approach
 
#include <bits/stdc++.h>
using namespace std;
 
// Function to return the index of 2*K-A[i]
int findElement(int A[], int low,
                int high, int key)
{
    int ans = -1;
 
    // Binary search
    while (low <= high) {
        int mid = low + (high - low) / 2;
        if (key <= A[mid]) {
            ans = mid;
            high = mid - 1;
        }
        else
            low = mid + 1;
    }
    return ans;
}
 
// Count the number of pairs
int countPairs(int A[], int& N, int& k)
{
    sort(A, A + N);
    int count = 0;
 
    // Loop to count the number of pairs
    for (int i = 0; i < N; i++) {
        int index
            = findElement(A, i + 1, N - 1,
                          2 * k - A[i]);
        if (index != -1)
            count += N - index;
    }
    return count;
}
 
// Driver Code
int main()
{
    int A[] = { 5, 1, 3, 4 };
    int N = sizeof(A) / sizeof(A[0]);
    int K = 3;
 
    // Function call
    cout << countPairs(A, N, K);
    return 0;
}

Java

// Java code to implement the approach
import java.util.*;
 
class GFG {    
 
  // Function to return the index of 2*K-A[i]
  static int findElement(int A[], int low,
                         int high, int key)
  {
    int ans = -1;
 
    // Binary search
    while (low <= high) {
      int mid = low + (high - low) / 2;
      if (key <= A[mid]) {
        ans = mid;
        high = mid - 1;
      }
      else
        low = mid + 1;
    }
    return ans;
  }
 
  // Count the number of pairs
  static int countPairs(int A[], int N, int k)
  {
    Arrays.sort(A);
    int count = 0;
 
    // Loop to count the number of pairs
    for (int i = 0; i < N; i++) {
      int index
        = findElement(A, i + 1, N - 1,
                      2 * k - A[i]);
      if (index != -1)
        count += N - index;
    }
    return count;
  }
 
  // Driver Code
  public static void main (String[] args) {
    int A[] = { 5, 1, 3, 4 };
    int N = A.length;
    int K = 3;
 
    // Function call
    System.out.print(countPairs(A, N, K));
  }
}
 
// This code is contributed by hrithikgarg03188.

Python3

# Python3 program for above approach
 
# Function to return the index of 2*K-A[i]
def findElement(A, low, high, key):
    ans = -1
 
    # binary search
    while (low <= high):
        mid = low + (high - low)//2
        if key <= A[mid]:
            ans = mid
            high = mid - 1
        else:
            low = mid + 1
    return ans
 
# Count the number of pairs
def countPairs(A, N, k):
    A.sort()
    count = 0
     
    # Loop to count the number of pairs
    for i in range(N):
        index = findElement(A, i + 1, N - 1, 2 * k - A[i])
        if index != -1:
            count += N - index
    return count
 
# Driver code
A = [5, 1, 3, 4]
N = len(A)
K = 3
 
# Function call
print(countPairs(A, N, K))
 
# this code is contributed by phasing17

C#

// C# code to implement the approach
using System;
 
public class GFG {
    // Function to return the index of 2*K-A[i]
    static int findElement(int[] A, int low, int high,
                           int key)
    {
        int ans = -1;
 
        // Binary search
        while (low <= high) {
            int mid = low + (high - low) / 2;
            if (key <= A[mid]) {
                ans = mid;
                high = mid - 1;
            }
            else
                low = mid + 1;
        }
        return ans;
    }
 
    // Count the number of pairs
    static int countPairs(int[] A, int N, int k)
    {
        Array.Sort(A);
        int count = 0;
 
        // Loop to count the number of pairs
        for (int i = 0; i < N; i++) {
            int index = findElement(A, i + 1, N - 1,
                                    2 * k - A[i]);
            if (index != -1)
                count += N - index;
        }
        return count;
    }
 
    // Driver Code
    static public void Main()
    {
        int[] A = { 5, 1, 3, 4 };
        int N = A.Length;
        int K = 3;
 
        // Function call
        Console.Write(countPairs(A, N, K));
    }
}
 
// This code is contributed by Rohit Pradhan

Javascript

<script>
      // JavaScript code for the above approach
 
      // Function to return the index of 2*K-A[i]
      function findElement(A, low,
          high, key) {
          let ans = -1;
 
          // Binary search
          while (low <= high) {
              let mid = low + Math.floor((high - low) / 2);
              if (key <= A[mid]) {
                  ans = mid;
                  high = mid - 1;
              }
              else
                  low = mid + 1;
          }
          return ans;
      }
 
      // Count the number of pairs
      function countPairs(A, N, k) {
          A.sort(function (a, b) { return a - b });
          let count = 0;
 
          // Loop to count the number of pairs
          for (let i = 0; i < N; i++) {
              let index
                  = findElement(A, i + 1, N - 1,
                      2 * k - A[i]);
              if (index != -1)
                  count += N - index;
          }
          return count;
      }
 
      // Driver Code
      let A = [5, 1, 3, 4];
      let N = A.length;
      let K = 3;
 
      // Function call
      document.write(countPairs(A, N, K));
 
  // This code is contributed by Potta Lokesh
  </script>

Output

Time Complexity: O(N log N)
Auxiliary Space: O(N)

Another Approach:

Sort the array A[] in non-decreasing order using std::sort() function.
Initialize a count variable to zero which will keep track of the number of pairs that have an average greater than or equal to K.
Loop through the array A[] from left to right.
For each element A[i], calculate the target value using the formula 2K – A[i]. This is because for two numbers X and Y, their average is greater than or equal to K if and only if (X + Y)/2 >= K, which simplifies to X + Y >= 2K.
Use binary search to find the index of the first element in the array A[] that is greater than or equal to the target value. Initialize left to i + 1 and right to N – 1, and repeat until left > right. Inside the loop, calculate mid as the average of left and right, and compare A[mid] with the target value. If A[mid] is greater than or equal to the target, set right to mid – 1, otherwise set left to mid + 1.
After the binary search, the index variable should contain the index of the first element in A[] that is greater than or equal to the target value. If index is less than N, then there are N – index pairs (A[i], A[j]), where i < j and A[j] >= target, that have an average greater than or equal to K. Add this count to the overall count variable.
After the loop, return the count variable as the final answer.

Example:

C++

// C++ code to implement the approach
 
#include <bits/stdc++.h>
using namespace std;
 
// Function to return the index of 2*K-A[i]
int countPairs(int A[], int& N, int& K)
{
sort(A, A + N);
int count = 0;
  // Loop to count the number of pairs
for (int i = 0; i < N; i++) {
    int target = 2*K - A[i];
    int left = i + 1, right = N - 1;
    while (left <= right) {
        int mid = left + (right - left) / 2;
        if (A[mid] >= target) {
            right = mid - 1;
        } else {
            left = mid + 1;
        }
    }
    int index = left;
    if (index != N)
        count += N - index;
}
return count;
}
 
// Driver Code
int main()
{
    int A[] = { 5, 1, 3, 4 };
    int N = sizeof(A) / sizeof(A[0]);
    int K = 3;
 
    // Function call
    cout << countPairs(A, N, K);
    return 0;
}

C#

// C# code to implement the approach
using System;
 
public class Program {
    // Function to return the index of 2*K-A[i]
    public static int CountPairs(int[] A, int N, int K) {
        Array.Sort(A);
        int count = 0;
        // Loop to count the number of pairs
        for (int i = 0; i < N; i++) {
            int target = 2 * K - A[i];
            int left = i + 1, right = N - 1;
            while (left <= right) {
                int mid = left + (right - left) / 2;
                if (A[mid] >= target) {
                    right = mid - 1;
                } else {
                    left = mid + 1;
                }
            }
            int index = left;
            if (index != N)
                count += N - index;
        }
        return count;
    }
 
    public static void Main() {
        int[] A = { 5, 1, 3, 4 };
        int N = A.Length;
        int K = 3;
        int result = CountPairs(A, N, K);
        Console.WriteLine(result);
    }
}

Python3

# Python code to implement the approach
 
def countPairs(A, N, K):
    A.sort()
    count = 0
    # Loop to count the number of pairs
    for i in range(N):
        target = 2 * K - A[i]
        left, right = i + 1, N - 1
        while left <= right:
            mid = left + (right - left) // 2
            if A[mid] >= target:
                right = mid - 1
            else:
                left = mid + 1
        index = left
        if index != N:
            count += N - index
    return count
 
# Driver code
A = [5, 1, 3, 4]
N = len(A)
K = 3
result = countPairs(A, N, K)
print(result)

Java

// Java code to implement the approach
 
import java.util.Arrays;
 
class Main {
    // Function to return the index of 2*K-A[i]
    static int countPairs(int[] A, int N, int K)
    {
        Arrays.sort(A);
        int count = 0;
        // Loop to count the number of pairs
        for (int i = 0; i < N; i++) {
            int target = 2 * K - A[i];
            int left = i + 1, right = N - 1;
            while (left <= right) {
                int mid = left + (right - left) / 2;
                if (A[mid] >= target) {
                    right = mid - 1;
                }
                else {
                    left = mid + 1;
                }
            }
            int index = left;
            if (index != N)
                count += N - index;
        }
        return count;
    }
 
    // Driver code
    public static void main(String[] args)
    {
        int[] A = { 5, 1, 3, 4 };
        int N = A.length;
        int K = 3;
        // Function call
        System.out.println(countPairs(A, N, K));
    }
}

Javascript

function countPairs(A, N, K) {
  A.sort(function(a, b) { return a - b; });
  let count = 0;
  for (let i = 0; i < N; i++) {
    let target = 2 * K - A[i];
    let left = i + 1, right = N - 1;
    while (left <= right) {
      let mid = Math.floor((left + right) / 2);
      if (A[mid] >= target) {
        right = mid - 1;
      } else {
        left = mid + 1;
      }
    }
    let index = left;
    if (index != N)
      count += N - index;
  }
  return count;
}
 
let A = [5, 1, 3, 4];
let N = A.length;
let K = 3;
let result = countPairs(A, N, K);
console.log(result);

Output

Time Complexity: O(N*logN)

Auxiliary Space: O(1)

Count of subarrays of size K with average at least M

prophet1999

Improve

Article Tags :

Practice Tags :

Count of pairs of Array elements with average at least K

C++14

Java

Python3

C#

Javascript

C++

C#

Python3

Java

Javascript

Similar Reads

Thank You!

What kind of Experience do you want to share?