Count subarrays made up of elements having exactly K set bits

Given an array arr[] consisting of N integers and an integer K, the task is to count the number of subarrays possible consisting of elements having exactly K set bits.

Examples:

Input: arr[] = {4, 2, 1, 5, 6}, K = 2
Output: 3
Explanation: The subarrays made up of elements having exactly 2 set bits are {5}, {6} and {5, 6}.

Input: arr[] = {4, 2, 1, 5, 6}, K = 1
Output: 6

Naive Approach: The simplest approach to solve the problem is to generate all possible subarrays of the given array and count those subarrays made up of elements having exactly K set bits. Finally, print the count of such subarrays. 

Time Complexity: O(N³log(M)), where M is the largest element in the array.
Auxiliary Space: O(1)

Efficient Approach: The idea is to keep the track of consecutive array elements with K set bits and find the count of subarrays with those consecutive sets of elements. Follow the steps below to solve the problem:

• Initialize a variable, say res as 0, to store the total count of subarrays consisting of elements having K set bits. Initialize a variable, say count as 0, to store the count of a consecutive set of elements having K set bits.
• Traverse the given array arr[] and perform the following steps:
  • If the current element arr[i] has K set bits, then increment count by 1.
  • Otherwise, increment the value of res by (count*(count - 1)) / 2 as the total number of subarrays formed by the previous consecutive elements and set count to 0.
• After completing the above steps, print the value of res as the resultant count of subarrays.

Below is the implementation of the above approach:

// C++ program for the above approach
#include <iostream>
using namespace std;

// Function to count the number
// of set bits in an integer N
int countSet(int N)
{

  // Stores the count of set bits
  int ans = 0;

  // While N is non-zero
  while(N)
  {

    // If the LSB is 1, then
    // increment ans by 1
    ans += N & 1;
    N >>= 1;
  }

  // Return the total set bits
  return ans;
}

// Function to count the number of
// subarrays having made up of
// elements having K set bits
int countSub(int *arr,int k)
{

  // Stores the total count of
  // resultant subarrays
  int ans = 0;
  int setK = 0;

  // Traverse the given array
  for(int i = 0; i < 5; i++)
  {

    // If the current element
    // has K set bits
    if(countSet(arr[i]) == k)
      setK += 1;

    // Otherwise
    else
      setK = 0;

    // Increment count of subarrays
    ans += setK;
  }

  // Return total count of subarrays
  return ans;
}

// Driver Code
int main() 
{
  int arr[] = {4, 2, 1, 5, 6};
  int K = 2;

  // Function Call
  cout<<(countSub(arr, K));
  return 0;
}

// This code is contributed by rohitsingh07052.

// Java program for the above approach
import java.util.*;

class GFG{

// Function to count the number
// of set bits in an integer N
static int countSet(int N)
{
    
    // Stores the count of set bits
    int ans = 0;
    
    // While N is non-zero
    while(N > 0)
    {
        
        // If the LSB is 1, then
        // increment ans by 1
        ans += N & 1;
        N >>= 1;
    }
    
    // Return the total set bits
    return ans;
}

// Function to count the number of
// subarrays having made up of
// elements having K set bits
static int countSub(int []arr,int k)
{
    
    // Stores the total count of
    // resultant subarrays
    int ans = 0;
    int setK = 0;
    
    // Traverse the given array
    for(int i = 0; i < 5; i++)
    {
    
        // If the current element
        // has K set bits
        if (countSet(arr[i]) == k)
            setK += 1;
        
        // Otherwise
        else
            setK = 0;
        
        // Increment count of subarrays
        ans += setK;
    }
    
    // Return total count of subarrays
    return ans;
}

// Driver Code
public static void main(String[] args) 
{
    int arr[] = {4, 2, 1, 5, 6};
    int K = 2;
    
    // Function Call
    System.out.print(countSub(arr, K));
}
}

// This code is contributed by 29AjayKumar 

# Python3 program for the above approach

# Function to count the number
# of set bits in an integer N
def countSet(N):

    # Stores the count of set bits
    ans = 0

    # While N is non-zero
    while N:

        # If the LSB is 1, then
        # increment ans by 1
        ans += N & 1
        N >>= 1

    # Return the total set bits
    return ans

# Function to count the number of
# subarrays having made up of
# elements having K set bits
def countSub(arr, k):

    # Stores the total count of
    # resultant subarrays
    ans = 0
    setK = 0

    # Traverse the given array
    for i in arr:

        # If the current element
        # has K set bits
        if countSet(i) == k:
            setK += 1

        # Otherwise
        else:
            setK = 0

        # Increment count of subarrays
        ans += setK

    # Return total count of subarrays
    return ans


# Driver Code
arr = [4, 2, 1, 5, 6]
K = 2

# Function Call
print(countSub(arr, K))

// C# program for the above approach
using System;

class GFG {

  // Function to count the number
  // of set bits in an integer N
  static int countSet(int N)
  {

    // Stores the count of set bits
    int ans = 0;

    // While N is non-zero
    while (N > 0) 
    {

      // If the LSB is 1, then
      // increment ans by 1
      ans += N & 1;
      N >>= 1;
    }

    // Return the total set bits
    return ans;
  }

  // Function to count the number of
  // subarrays having made up of
  // elements having K set bits
  static int countSub(int[] arr, int k)
  {

    // Stores the total count of
    // resultant subarrays
    int ans = 0;
    int setK = 0;

    // Traverse the given array
    for (int i = 0; i < 5; i++) {

      // If the current element
      // has K set bits
      if (countSet(arr[i]) == k)
        setK += 1;

      // Otherwise
      else
        setK = 0;

      // Increment count of subarrays
      ans += setK;
    }

    // Return total count of subarrays
    return ans;
  }

  // Driver Code
  public static void Main(string[] args)
  {
    int[] arr = { 4, 2, 1, 5, 6 };
    int K = 2;

    // Function Call
    Console.WriteLine(countSub(arr, K));
  }
}

// This code is contributed by ukasp.

<script>

// JavaScript program for the above approach

// Function to count the number
// of set bits in an integer N
function countSet(N)
{

// Stores the count of set bits
let ans = 0;

// While N is non-zero
while(N)
{

    // If the LSB is 1, then
    // increment ans by 1
    ans += N & 1;
    N >>= 1;
}

// Return the total set bits
return ans;
}

// Function to count the number of
// subarrays having made up of
// elements having K set bits
function countSub(arr,k)
{

// Stores the total count of
// resultant subarrays
let ans = 0;
let setK = 0;

// Traverse the given array
for(let i = 0; i < 5; i++)
{

    // If the current element
    // has K set bits
    if(countSet(arr[i]) == k)
    setK += 1;

    // Otherwise
    else
    setK = 0;

    // Increment count of subarrays
    ans += setK;
}

// Return total count of subarrays
return ans;
}

// Driver Code

let arr = [4, 2, 1, 5, 6];
let K = 2;

// Function Call
document.write((countSub(arr, K)));

// This code is contributed by Mayank Tyagi
</script>

3

Time Complexity: O(N*log(M)), where M is the largest element in the array.
Auxiliary Space: O(1)