Count subarrays made up of elements having exactly K set bits

Last Updated : 31 Mar, 2021

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

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

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

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

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

Javascript

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