Skip to content
Related Articles

Related Articles

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(N3log(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>
Output: 
3

 

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

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.

My Personal Notes arrow_drop_up
Recommended Articles
Page :