Number of ways to divide an array into K equal sum sub-arrays

Given an integer K and an array arr[] of N integers, the task is to find the number of ways to split the array into K equal sum sub-arrays of non-zero lengths.

Examples:

Input: arr[] = {0, 0, 0, 0}, K = 3
Output: 3
All possible ways are:
{{0}, {0}, {0, 0}}
{{0}, {0, 0}, {0}}
{{0, 0}, {0}, {0}}



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

Approach: This problem can be solved using dynamic programming. Following will be our algorithm:

  1. Find the sum of all the elements of the array and store it in a variable SUM.
    Before going to step 2, let’s try and understand the states of the DP.
    For that, visualize putting bars to divide array into K equal parts. So, we have to put K – 1 bars in total.
    Thus, our states of dp will contain 2 terms.

    • i – index of the element we are currently upon.
    • ck – number of bars we have already inserted + 1.

    dp[i][ck] can be defined as number of ways to put the remaining bars into the array such that it gets divided into K equal halves. Now, lets proceed to step 2 of our algorithm.

  2. Call a recursive function with i = 0 and ck = 1 and the recurrence relation will be:

    Case 1: sum upto index i equals ((SUM)/k)* ck
    dp[i][ck] = dp[i+1][ck] + dp[i+1][ck+1]
    Case 2: sum upto index not i equals ((SUM)/k)* ck
    dp[i][ck] = dp[i+1][ck]

Below is the implementation of the above approach:

C++

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ implementation of the approach
#include <bits/stdc++.h>
#define max_size 20
#define max_k 20
using namespace std;
  
// Array to store the states of DP
int dp[max_size][max_k];
  
// Array to check if a
// state has been solved before
bool v[max_size][max_k];
  
// To store the sum of
// the array elements
int sum = 0;
  
// Function to find the sum of
// all the array elements
void findSum(int arr[], int n)
{
    for (int i = 0; i < n; i++)
        sum += arr[i];
}
  
// Function to return the number of ways
int cntWays(int arr[], int i, int ck,
            int k, int n, int curr_sum)
{
    // If sum is not divisible by k
    // answer will be zero
    if (sum % k != 0)
        return 0;
    if (i != n and ck == k + 1)
        return 0;
  
    // Base case
    if (i == n) {
        if (ck == k + 1)
            return 1;
        else
            return 0;
    }
  
    // To check if a state
    // has been solved before
    if (v[i][ck])
        return dp[i][ck];
  
    // Sum of all the numbers from the beginning
    // of the array
    curr_sum += arr[i];
  
    // Setting the current state as solved
    v[i][ck] = 1;
  
    // Recurrence relation
    dp[i][ck] = cntWays(arr, i + 1, ck, k, n, curr_sum);
    if (curr_sum == (sum / k) * ck)
        dp[i][ck] += cntWays(arr, i + 1, ck + 1, k, n, curr_sum);
  
    // Returning solved state
    return dp[i][ck];
}
  
// Driver code
int main()
{
    int arr[] = { 1, -1, 1, -1, 1, -1 };
    int n = sizeof(arr) / sizeof(int);
    int k = 2;
  
    // Function call to find the
    // sum of the array elements
    findSum(arr, n);
  
    // Print the number of ways
    cout << cntWays(arr, 0, 1, k, n, 0);
}

chevron_right


Java

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java implementation of the approach
class GFG 
{
      
static int max_size= 20;
static int max_k =20;
  
// Array to store the states of DP
static int [][]dp = new int[max_size][max_k];
  
// Array to check if a
// state has been solved before
static boolean [][]v = new boolean[max_size][max_k];
  
// To store the sum of
// the array elements
static int sum = 0;
  
// Function to find the sum of
// all the array elements
static void findSum(int arr[], int n)
{
    for (int i = 0; i < n; i++)
        sum += arr[i];
}
  
// Function to return the number of ways
static int cntWays(int arr[], int i, int ck,
            int k, int n, int curr_sum)
{
    // If sum is not divisible by k
    // answer will be zero
    if (sum % k != 0)
        return 0;
    if (i != n && ck == k + 1)
        return 0;
  
    // Base case
    if (i == n) 
    {
        if (ck == k + 1)
            return 1;
        else
            return 0;
    }
  
    // To check if a state
    // has been solved before
    if (v[i][ck])
        return dp[i][ck];
  
    // Sum of all the numbers from the beginning
    // of the array
    curr_sum += arr[i];
  
    // Setting the current state as solved
    v[i][ck] = true;
  
    // Recurrence relation
    dp[i][ck] = cntWays(arr, i + 1, ck, k, n, curr_sum);
    if (curr_sum == (sum / k) * ck)
        dp[i][ck] += cntWays(arr, i + 1, ck + 1, k, n, curr_sum);
  
    // Returning solved state
    return dp[i][ck];
}
  
// Driver code
public static void main(String[] args)
{
    int arr[] = { 1, -1, 1, -1, 1, -1 };
    int n = arr.length;
    int k = 2;
  
    // Function call to find the
    // sum of the array elements
    findSum(arr, n);
  
    // Print the number of ways
    System.out.println(cntWays(arr, 0, 1, k, n, 0));
}
}
  
// This code is contributed by Princi Singh

chevron_right


Python3

filter_none

edit
close

play_arrow

link
brightness_4
code

# Python3 implementation of the approach
import numpy as np
  
max_size = 20 
max_k = 20 
  
  
# Array to store the states of DP 
dp = np.zeros((max_size,max_k)); 
  
# Array to check if a 
# state has been solved before 
v = np.zeros((max_size,max_k)); 
  
# To store the sum of 
# the array elements 
sum = 0
  
# Function to find the sum of 
# all the array elements 
def findSum(arr, n) : 
    global sum
    for i in range(n) :
        sum += arr[i]; 
  
  
# Function to return the number of ways 
def cntWays(arr, i, ck, k, n,  curr_sum) : 
  
    # If sum is not divisible by k 
    # answer will be zero 
    if (sum % k != 0) :
        return 0
    if (i != n and ck == k + 1) :
        return 0
  
    # Base case 
    if (i == n) :
        if (ck == k + 1) :
            return 1
        else :
            return 0
  
    # To check if a state 
    # has been solved before 
    if (v[i][ck]) :
        return dp[i][ck]; 
  
    # Sum of all the numbers from the beginning 
    # of the array 
    curr_sum += arr[i]; 
  
    # Setting the current state as solved 
    v[i][ck] = 1
  
    # Recurrence relation 
    dp[i][ck] = cntWays(arr, i + 1, ck, k, n, curr_sum); 
    if (curr_sum == (sum / k) * ck)  :
        dp[i][ck] += cntWays(arr, i + 1, ck + 1, k, n, curr_sum); 
  
    # Returning solved state 
    return dp[i][ck]; 
   
  
# Driver code 
if __name__ == "__main__"
  
    arr = [ 1, -1, 1, -1, 1, -1 ]; 
    n = len(arr); 
    k = 2
  
    # Function call to find the 
    # sum of the array elements 
    findSum(arr, n); 
  
    # Print the number of ways 
    print(cntWays(arr, 0, 1, k, n, 0)); 
  
    # This code is contributed by AnkitRai01

chevron_right


C#

filter_none

edit
close

play_arrow

link
brightness_4
code

// C# implementation of the approach
using System;     
      
class GFG 
{
      
static int max_size= 20;
static int max_k =20;
  
// Array to store the states of DP
static int [,]dp = new int[max_size, max_k];
  
// Array to check if a
// state has been solved before
static Boolean [,]v = new Boolean[max_size, max_k];
  
// To store the sum of
// the array elements
static int sum = 0;
  
// Function to find the sum of
// all the array elements
static void findSum(int []arr, int n)
{
    for (int i = 0; i < n; i++)
        sum += arr[i];
}
  
// Function to return the number of ways
static int cntWays(int []arr, int i, int ck,
            int k, int n, int curr_sum)
{
    // If sum is not divisible by k
    // answer will be zero
    if (sum % k != 0)
        return 0;
    if (i != n && ck == k + 1)
        return 0;
  
    // Base case
    if (i == n) 
    {
        if (ck == k + 1)
            return 1;
        else
            return 0;
    }
  
    // To check if a state
    // has been solved before
    if (v[i, ck])
        return dp[i, ck];
  
    // Sum of all the numbers from the beginning
    // of the array
    curr_sum += arr[i];
  
    // Setting the current state as solved
    v[i, ck] = true;
  
    // Recurrence relation
    dp[i,ck] = cntWays(arr, i + 1, ck, k, n, curr_sum);
    if (curr_sum == (sum / k) * ck)
        dp[i, ck] += cntWays(arr, i + 1, ck + 1, k, n, curr_sum);
  
    // Returning solved state
    return dp[i, ck];
}
  
// Driver code
public static void Main(String[] args)
{
    int []arr = { 1, -1, 1, -1, 1, -1 };
    int n = arr.Length;
    int k = 2;
  
    // Function call to find the
    // sum of the array elements
    findSum(arr, n);
  
    // Print the number of ways
    Console.WriteLine(cntWays(arr, 0, 1, k, n, 0));
}
}
  
// This code contributed by Rajput-Ji

chevron_right


Output:

2

Time Complexity: O(n*k)



My Personal Notes arrow_drop_up

Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.