Check if an array can be split into K non-overlapping subarrays whose Bitwise AND values are equal

Given an array arr[] of size N and a positive integer K, the task is to check if the array can be split into K non-overlapping and non-empty subarrays such that Bitwise AND of all the subarrays are equal. If found to be true, then print “YES”. Otherwise, print “NO”.

Examples:

Input: arr[] = { 3, 2, 2, 6, 2 }, K = 3 
Output: YES 
Explanation:
Splitting the array into K( = 3) subarrays as { { 3, 2 }, { 2, 6 }, { 2 } } 
Therefore, the required output is YES.

Input: arr[] = { 4, 3, 5, 2 }, K = 3 
Output: NO

Naive Approach: The simplest approach to solve this problem is to split the array into K subarrays in all possible ways and in every possible way, check if Bitwise AND of all K subarrays are equal or not. If found to be true for any split, then print “YES”. Otherwise, print “NO”. 

Time Complexity: O(N³)
Auxiliary Space: O(1)

Efficient Approach: To optimize the above approach, the idea is to use the fact that if i^th bit of at least one element of the subarray is 0, then ith bit of bitwise AND of that subarray is also 0. Follow the steps below to solve the problem:

• Initialize a variable, say flag, to check if the array can be split into K subarrays such that Bitwise AND of all the subarrays are equal.
• Initialize a 2D array, say pref[][], where pref[i][j] stores the count of contiguous array elements up to i^th index whose j^th bit is set.
• Iterate over the range [0, N – K]. For every j^th bit of i^th element, check for the following conditions: 
  • If j^th bit of all the array elements up to i^th index are set and j^th bit of at least one element of the array after i^th index are not set, then update flag = false.
  • If j^th bit of at least one array elements up to i^th index are not set and j^th bit of all the array elements after the i^th index are set then update flag = false
• Finally, check if flag equal to true or not. If found to be true, then print “YES”.
• Otherwise, print “NO”.

Below is the implementation of the above approach:

C++

// C++ program to implement
// the above approach

#include <bits/stdc++.h>
using namespace std;

// Utility function to check if the array
// can be split into K subarrays whose
// bitwise AND are equal
bool equalPartitionUtil(int arr[], int N, int K)
{

    // pref[i][j]: Stores count of contiguous
    // array elements upto i-th index whose
    // j-th bit is set
    int pref[N][32];

    // Initialize pref[][] array
    memset(pref, 0, sizeof(pref));

    // Fill the prefix array
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < 32; j++) {
            if (i) {

                // Check if j-th bit set or not
                int X = ((arr[i] & (1 << j)) > 0);

                // Update pref[i][j]
                pref[i][j] = pref[i - 1][j] + X;
            }

            else {

                // Update pref[i][j]
                pref[i][j] = ((arr[i] & (1 << j)) > 0);
            }
        }
    }

    // Iterate over the range[0, N - K]
    for (int i = 0; i < N - K + 1; i++) {
        bool flag = true;

        for (int j = 0; j < 32; j++) {

            // Get count of elements that have
            // jth bit set
            int cnt = pref[i][j];

            // Check if first case is satisfied
            if (cnt == i + 1
                && pref[N - 1][j] - pref[i][j] != N - i - 1)
                flag = false;

            // Check if second case is satisfied
            if (cnt != i + 1
                && N - i - 1 - (pref[N - 1][j] - pref[i][j]) < K - 1)
                flag = false;
        }

        if (flag)
            return true;
    }

    return false;
}

// Function to check if the array
// can be split into K subarrays
// having equal value of bitwise AND
void equalPartition(int arr[], int N, int K)
{
    if (equalPartitionUtil(arr, N, K))
        cout << "YES";
    else
        cout << "NO";
}

// Driver code
int main()
{
    // Given array
    int arr[] = { 3, 2, 2, 6, 2 };

    // Size of the array
    int N = sizeof(arr) / sizeof(arr[0]);

    // Given K
    int K = 3;

    // Function Call
    equalPartition(arr, N, K);

    return 0;
}

Java

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

class GFG{

// Utility function to check if the array
// can be split into K subarrays whose
// bitwise AND are equal
static boolean equalPartitionUtil(int arr[],
                                  int N, int K)
{

    // pref[i][j]: Stores count of contiguous
    // array elements upto i-th index whose
    // j-th bit is set
    int [][]pref = new int[N][32];

    // Fill the prefix array
    for(int i = 0; i < N; i++)
    {
        for(int j = 0; j < 32; j++)
        {
            if (i > 0)
            {

                // Check if j-th bit set or not
                int X = ((arr[i] & (1 << j)) > 0) ? 1 : 0;

                // Update pref[i][j]
                pref[i][j] = pref[i - 1][j] + X;
            }

            else
            {

                // Update pref[i][j]
                pref[i][j] = ((arr[i] & (1 << j)) > 0) ? 1 : 0;
            }
        }
    }

    // Iterate over the range[0, N - K]
    for(int i = 0; i < N - K + 1; i++)
    {
        boolean flag = true;

        for(int j = 0; j < 32; j++)
        {

            // Get count of elements that have
            // jth bit set
            int cnt = pref[i][j];

            // Check if first case is satisfied
            if (cnt == i + 1 && pref[N - 1][j] -
                   pref[i][j] != N - i - 1)
                flag = false;

            // Check if second case is satisfied
            if (cnt != i + 1 && N - i - 1 - (
                pref[N - 1][j] - pref[i][j]) < K - 1)
                flag = false;
        }
        if (flag)
            return true;
    }
    return false;
}

// Function to check if the array
// can be split into K subarrays
// having equal value of bitwise AND
static void equalPartition(int arr[], int N, int K)
{
    if (equalPartitionUtil(arr, N, K))
        System.out.print("YES");
    else
        System.out.print("NO");
}

// Driver code
public static void main(String[] args)
{

    // Given array
    int arr[] = { 3, 2, 2, 6, 2 };

    // Size of the array
    int N = arr.length;

    // Given K
    int K = 3;

    // Function Call
    equalPartition(arr, N, K);
}
}

// This code is contributed by shikhasingrajput

Python3

# Python3 program to implement
# the above approach

# Utility function to check if the array
# can be split into K subarrays whose
# bitwise AND are equal
def equalPartitionUtil(arr, N, K):

    # pref[i][j]: Stores count of contiguous
    # array elements upto i-th index whose
    # j-th bit is set
    pref = [[0 for x in range(32)]for y in range(N)]

    # Fill the prefix array
    for i in range(N):
        for j in range(32):
            if (i):

                # Check if j-th bit set or not
                X = ((arr[i] & (1 << j)) > 0)

                # Update pref[i][j]
                pref[i][j] = pref[i - 1][j] + X

            else:

                # Update pref[i][j]
                pref[i][j] = ((arr[i] & (1 << j)) > 0)

    # Iterate over the range[0, N - K]
    for i in range(N - K + 1):
        flag = True
        for j in range(32):

            # Get count of elements that have
            # jth bit set
            cnt = pref[i][j]

            # Check if first case is satisfied
            if (cnt == i + 1
                    and pref[N - 1][j] - pref[i][j] != N - i - 1):
                flag = False

            # Check if second case is satisfied
            if (cnt != i + 1
                    and N - i - 1 - (pref[N - 1][j] - pref[i][j]) < K - 1):
                flag = False
        if (flag):
            return True
    return False

# Function to check if the array
# can be split into K subarrays
# having equal value of bitwise AND
def equalPartition(arr, N, K):
    if (equalPartitionUtil(arr, N, K)):
        print("YES")
    else:
        print("NO")

# Driver code
if __name__ == "__main__":

    # Given array
    arr = [3, 2, 2, 6, 2]

    # Size of the array
    N = len(arr)

    # Given K
    K = 3

    # Function Call
    equalPartition(arr, N, K)

    # This code is contributed by chitranayal.

C#

// C# program to implement
// the above approach
using System;

class GFG{

// Utility function to check if the array
// can be split into K subarrays whose
// bitwise AND are equal
static bool equalPartitionUtil(int []arr,
                               int N, int K)
{

    // pref[i,j]: Stores count of contiguous
    // array elements upto i-th index whose
    // j-th bit is set
    int [,]pref = new int[N, 32];

    // Fill the prefix array
    for(int i = 0; i < N; i++)
    {
        for(int j = 0; j < 32; j++)
        {
            if (i > 0)
            {

                // Check if j-th bit set or not
                int X = ((arr[i] & (1 << j)) > 0) ? 1 : 0;

                // Update pref[i,j]
                pref[i, j] = pref[i - 1, j] + X;
            }

            else
            {

                // Update pref[i,j]
                pref[i, j] = ((arr[i] & (1 << j)) > 0) ? 1 : 0;
            }
        }
    }

    // Iterate over the range[0, N - K]
    for(int i = 0; i < N - K + 1; i++)
    {
        bool flag = true;

        for(int j = 0; j < 32; j++)
        {

            // Get count of elements that have
            // jth bit set
            int cnt = pref[i, j];

            // Check if first case is satisfied
            if (cnt == i + 1 && pref[N - 1, j] -
                pref[i, j] != N - i - 1)
                flag = false;

            // Check if second case is satisfied
            if (cnt != i + 1 && N - i - 1 - (
                pref[N - 1, j] - pref[i, j]) < K - 1)
                flag = false;
        }
        if (flag)
            return true;
    }
    return false;
}

// Function to check if the array
// can be split into K subarrays
// having equal value of bitwise AND
static void equalPartition(int []arr, int N, int K)
{
    if (equalPartitionUtil(arr, N, K))
        Console.Write("YES");
    else
        Console.Write("NO");
}

// Driver code
public static void Main(String[] args)
{

    // Given array
    int []arr = { 3, 2, 2, 6, 2 };

    // Size of the array
    int N = arr.Length;

    // Given K
    int K = 3;

    // Function Call
    equalPartition(arr, N, K);
}
}

// This code is contributed by shikhasingrajput

Javascript

<script>

// Javascript program to implement
// the above approach

// Utility function to check if the array
// can be split into K subarrays whose
// bitwise AND are equal
function equalPartitionUtil(arr, N, K)
{

    // pref[i][j]: Stores count of contiguous
    // array elements upto i-th index whose
    // j-th bit is set
    var pref = Array.from(Array(N), ()=> Array(32).fill(0));

    // Fill the prefix array
    for (var i = 0; i < N; i++) {
        for (var j = 0; j < 32; j++) {
            if (i) {

                // Check if j-th bit set or not
                var X = ((arr[i] & (1 << j)) > 0);

                // Update pref[i][j]
                pref[i][j] = pref[i - 1][j] + X;
            }

            else {

                // Update pref[i][j]
                pref[i][j] = ((arr[i] & (1 << j)) > 0);
            }
        }
    }

    // Iterate over the range[0, N - K]
    for (var i = 0; i < N - K + 1; i++) {
        var flag = true;

        for (var j = 0; j < 32; j++) {

            // Get count of elements that have
            // jth bit set
            var cnt = pref[i][j];

            // Check if first case is satisfied
            if (cnt == i + 1
                && pref[N - 1][j] - pref[i][j] != N - i - 1)
                flag = false;

            // Check if second case is satisfied
            if (cnt != i + 1
                && N - i - 1 - (pref[N - 1][j] - pref[i][j]) < K - 1)
                flag = false;
        }

        if (flag)
            return true;
    }

    return false;
}

// Function to check if the array
// can be split into K subarrays
// having equal value of bitwise AND
function equalPartition(arr, N, K)
{
    if (equalPartitionUtil(arr, N, K))
        document.write( "YES");
    else
        document.write( "NO");
}

// Driver code
// Given array
var arr = [ 3, 2, 2, 6, 2 ];

// Size of the array
var N = arr.length;

// Given K
var K = 3;

// Function Call
equalPartition(arr, N, K);

// This code is contributed by itsok.
</script>

Output:

YES

Time Complexity: O(32 * N)
Auxiliary Space: O(32 * N)