Count of pairs of Array elements with average at least K

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++ 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 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 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# 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

  <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>

4

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

Another Approach: 

1. Sort the array A[] in non-decreasing order using std::sort() function.
2. 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.
3. Loop through the array A[] from left to right.
4. 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.
5. 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.
6. 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.
7. After the loop, return the count variable as the final answer.

Example:

// 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# 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);
    }
}

# 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 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));
    }
}

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);

4

Time Complexity:  O(N*logN)

Auxiliary Space: O(1)

prophet1999

Follow

Improve

Article Tags :

• Sorting
• DSA
• Arrays
• Binary Search

Practice Tags :

• Arrays
• Binary Search
• Sorting