Count ways to partition Binary Array into subarrays containing K 0s each

Given a binary array arr[] of size N, and an integer K, the task is to calculate the number of ways to partition the array into non-overlapping subarrays, where each subarray has exactly K number 0s.

Examples:

Input: arr[] = [ 0, 0, 1, 1, 0, 1, 0], K = 2
Output: 3
Explanation: Different possible partitions are: 
{{0, 0}, {1, 1, 0, 1, 0}}, {{0, 0, 1}, {1, 0, 1, 0}}, {{0, 0, 1, 1}, {0, 1, 0}}. So, the output will be 3.

Input: arr[] = {0, 0, 1, 0, 1, 0}, K = 2
Output: 2

Input: arr[] = [1, 0, 1, 1], K = 2
Output: 0

Approach: The approach to solve the problem is based on the following idea:

If jth 0 is the last 0 for a subarray and (j+1)th 0 is the first 0 of another subarray, then the possible number of ways to partition into those two subarrays is one more than the number of 1s in between jth and (j+1)th 0.

From the above observation, it can be said that the total possible ways to partition the subarray is the multiplication of the count of 1s between K*x th and (K*x + 1)th 0, for all possible x such that K*x does not exceed the total count of 0s in the array.

Follow the illustration below for a better understanding of the problem,

Illustration:

Consider array arr[] = {0, 0, 1, 1, 0, 1, 0, 1, 0, 0}, K = 2

Index of 2nd 0 and 3rd 0 are 1 and 4
        => Total number of 1s in between = 2.
        => Possible partition with these 0s = 2 + 1 = 3.
        => Total possible partitions till now = 3

Index of 4th 0 and 5th 0 are 6 and 8
        => Total number of 1s in between = 1.
        => Possible partition with these 0s = 1 + 1 = 2.
        => Total possible partitions till now = 3*2 = 6

The possible partitions are 6: 
{{0, 0}, {1, 1, 0, 1, 0}, {1, 0, 0}}, {{0, 0}, {1, 1, 0, 1, 0, 1}, {0, 0}},  
{{0, 0, 1}, {1, 0, 1, 0}, {1, 0, 0}}, {{0, 0, 1}, {1, 0, 1, 0, 1}, {0, 0}}, 
{{0, 0, 1, 1}, {0, 1, 0}, {1, 0, 0}}, {{0, 0, 1, 1}, {0, 1, 0, 1}, {0, 0}} 

Follow the steps mentioned below to solve the problem:

• Initialize a counter variable to 1(claiming there exists at least one such possible way). 
• If there are less than K 0s or number of 0s is not divisible by K, then such partition is not possible.
• Then, for every possible (K*x)th and (K*x + 1)th number of 0, calculate the number of possible partitions using the above observation and multiply that with the counter variable to get the total possible partitions.
• Return the value of the counter variable.

Here is the code for the above approach:

// C++ program for above approach

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

// Function used to calculate the number of
// ways to divide the array
int number_of_ways(vector<int>& arr, int K)
{
    // Initialize a counter variable no_0 to
    // calculate the number of 0
    int no_0 = 0;

    // Initialize a vector to
    // store the indices of 0s
    vector<int> zeros;
    for (int i = 0; i < arr.size(); i++) {
        if (arr[i] == 0) {
            no_0++;
            zeros.push_back(i);
        }
    }

    // If number of 0 is not divisible by K
    // or no 0 in the sequence return 0
    if (no_0 % K || no_0 == 0)
        return 0;

    int res = 1;

    // For every (K*n)th and (K*n+1)th 0
    // calculate the distance between them
    for (int i = K; i < zeros.size();) {
        res *= (zeros[i] - zeros[i - 1]);
        i += K;
    }

    // Return the number of such partitions
    return res;
}

// Driver code
int main()
{
    vector<int> arr = { 0, 0, 1, 1, 0, 1, 0 };
    int K = 2;

    // Function call
    cout << number_of_ways(arr, K) << endl;
    return 0;
}

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

class GFG {
    // Function used to calculate the number of
    // ways to divide the array
    public static int number_of_ways(int arr[], int K)
    {
        // Initialize a counter variable no_0 to
        // calculate the number of 0
        int no_0 = 0;

        // Initialize a arraylist to
        // store the indices of 0s
        ArrayList<Integer> zeros = new ArrayList<Integer>();
        for (int i = 0; i < arr.length; i++) {
            if (arr[i] == 0) {
                no_0++;
                zeros.add(i);
            }
        }

        // If number of 0 is not divisible by K
        // or no 0 in the sequence return 0
        if ((no_0 % K != 0) || no_0 == 0)
            return 0;

        int res = 1;

        // For every (K*n)th and (K*n+1)th 0
        // calculate the distance between them
        for (int i = K; i < zeros.size();) {
            res *= (zeros.get(i) - zeros.get(i - 1));
            i += K;
        }

        // Return the number of such partitions
        return res;
    }

    // Driver Code
    public static void main(String[] args)
    {
        int arr[] = { 0, 0, 1, 1, 0, 1, 0 };
        int K = 2;

        // Function call
        System.out.println(number_of_ways(arr, K));
    }
}

// This code is contributed by Rohit Pradhan

# Python3 program for above approach

# Function used to calculate the number of
# ways to divide the array
def number_of_ways(arr, K):
  
    # Initialize a counter variable no_0 to
    # calculate the number of 0
    no_0 = 0
    
    # Initialize am array to
    # store the indices of 0s
    zeros = []
    for i in range(len(arr)):
        if arr[i] == 0:
            no_0 += 1
            zeros.append(i)
            
    # If number of 0 is not divisible by K
    # or no 0 in the sequence return 0
    if no_0 % K or no_0 == 0:
        return 0

    res = 1
    
    # For every (K*n)th and (K*n+1)th 0
    # calculate the distance between them
    i = K
    while (i < len(zeros)):
        res *= (zeros[i] - zeros[i - 1])
        i += K
        
    # Return the number of such partitions
    return res

# Driver code
arr = [0, 0, 1, 1, 0, 1, 0]
K = 2

# Function call
print(number_of_ways(arr, K))

# This code is contributed by phasing17.

// C# program for above approach

using System;
using System.Collections.Generic;

public class GFG 
{

  // Function used to calculate the number of
  // ways to divide the array
  public static int number_of_ways(int[] arr, int K)
  {

    // Initialize a counter variable no_0 to
    // calculate the number of 0
    int no_0 = 0;

    // Initialize a arraylist to
    // store the indices of 0s
    var zeros = new List<int>();
    for (int i = 0; i < arr.Length; i++) {
      if (arr[i] == 0) {
        no_0++;
        zeros.Add(i);
      }
    }

    // If number of 0 is not divisible by K
    // or no 0 in the sequence return 0
    if ((no_0 % K != 0) || no_0 == 0)
      return 0;

    int res = 1;

    // For every (K*n)th and (K*n+1)th 0
    // calculate the distance between them
    for (int i = K; i < zeros.Count;) {
      res *= (zeros[i] - zeros[i - 1]);
      i += K;
    }

    // Return the number of such partitions
    return res;
  }
  public static void Main(string[] args)
  {
    int[] arr = { 0, 0, 1, 1, 0, 1, 0 };
    int K = 2;

    // Function call
    Console.WriteLine(number_of_ways(arr, K));
  }
}

// this code was contributed by phasing17

    <script>
        // JavaScript program for above approach


        // Function used to calculate the number of
        // ways to divide the array
        const number_of_ways = (arr, K) => {
            // Initialize a counter variable no_0 to
            // calculate the number of 0
            let no_0 = 0;

            // Initialize a vector to
            // store the indices of 0s
            let zeros = [];
            for (let i = 0; i < arr.length; i++) {
                if (arr[i] == 0) {
                    no_0++;
                    zeros.push(i);
                }
            }

            // If number of 0 is not divisible by K
            // or no 0 in the sequence return 0
            if (no_0 % K || no_0 == 0)
                return 0;

            let res = 1;

            // For every (K*n)th and (K*n+1)th 0
            // calculate the distance between them
            for (let i = K; i < zeros.length;) {
                res *= (zeros[i] - zeros[i - 1]);
                i += K;
            }

            // Return the number of such partitions
            return res;
        }

        // Driver code

        let arr = [0, 0, 1, 1, 0, 1, 0];
        let K = 2;

        // Function call
        document.write(number_of_ways(arr, K));

    // This code is contributed by rakeshsahni

    </script>

3

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