Open In App

Partition of a set into k subsets with equal sum using BitMask and DP

Last Updated : 26 Nov, 2024
Comments
Improve
Suggest changes
Like Article
Like
Report

Given an integer array arr[] and an integer k, the task is to check if it is possible to divide the given array into k non-empty subsets of equal sum such that every array element is part of a single subset.

Examples:  

Input: arr[] = [2, 1, 4, 5, 6], k = 3 
Output: true
Explanation: Possible subsets of the given array are [2, 4], [1, 5] and [6]

Input: arr[] = [2, 1, 5, 5, 6], k = 3 
Output: false
Explanation: It is not possible to divide above array into 3 parts with equal sum.

Using Recursion – O(k*2^n) Time and O(n) Space

For the recursive approach, refer to partitioning a set into k subsets with equal sum. For each value of k, we use recursion to generate all possible subsequences of the array and check if any of them forms a valid subset with the required sum. Once a valid subset is found, its elements are marked as used, and the recursion continues for the remaining subsets. This process is repeated recursively for all k subsets, exploring all combinations until either the subsets are successfully formed or no further options remain.

Using Bitmasking and DP – O(n*2^n) and O(2^n) Space

The idea is to use mask to determine the current state. The current state will tell us about the subset already formed (which numbers are already selected). 
For example: arr[] = [2, 1, 4, 3, 5, 6, 2], mask = (1100101), which means that [2, 1, 5, 2] are already chosen in the current mask.
For any current state mask, the jth element will be added to it based on the following two conditions:

  • The jth bit is not set in the mask (mask & (1<<j) == 0)
  • sum (mask) + arr[j] <= target ( where target = (sum of array elements) / k)

Maintain a table dp[] such that dp[i] store the sum of elements in mask i. So, the dp transitions will be:
dp[i | (1 << j)] = (dp[i] + arr[j]) % target 

C++
// C++ program to check if the
// given array can be partitioned
// into K subsets with equal sum
// using bitmasking dp

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

bool isKPartitionPossible(vector<int> &arr, int k) {
    int n = arr.size();

    if (k == 1)

        // Return true as the entire array is the answer
        return true;

    // If total number of partitions exceeds
  	// size of the array
    if (n < k)
        return false;

    int sum = 0;
    for (int i = 0; i < n; i++)
        sum += arr[i];

    // If the array sum is not divisible by K
    if (sum % k != 0)

        // No such partitions are possible
        return false;

    // Required sum of each subset
    int target = sum / k;

    // Initialize dp vector with -1
    vector<int> dp(1 << n, -1);

    // Sum of empty subset is zero
    dp[0] = 0;

    // Iterate over all subsets/masks
    for (int mask = 0; mask < (1 << n); mask++) {

        // if current mask is invalid, continue
        if (dp[mask] == -1)
            continue;

        // Iterate over all array elements
        for (int i = 0; i < n; i++) {

            // Check if the current element can be added 
          	// to the current subset/mask
            if (!(mask & (1 << i)) && dp[mask] + arr[i] <= target) {
                dp[mask | (1 << i)] = (dp[mask] + arr[i]) % target;
            }
        }
    }

    // If the dp value of all elements used is zero, then
  	// partitioning is possible
    return dp[(1 << n) - 1] == 0;
}

int main() {
    vector<int> arr = {2, 1, 4, 5, 3, 3};
    int k = 2;

    if (isKPartitionPossible(arr, k)) {
        cout << "true";
    }
    else {
        cout << "false";
    }
}
Java
// Java program to check if the
// given array can be partitioned
// into K subsets with equal sum
// using bitmasking dp

import java.util.*;

class GfG {

    static boolean isKPartitionPossible(int[] arr, int k) {
        int n = arr.length;

        if (k == 1)
          
            // Return true as the entire array
          	// is the answer
            return true;

        // If total number of partitions exceeds size of the
        // array
        if (n < k)
            return false;

        int sum = 0;
        for (int i = 0; i < n; i++)
            sum += arr[i];

        // If the array sum is not divisible by K
        if (sum % k != 0)
          
            // No such partitions are possible
            return false;

        // Required sum of each subset
        int target = sum / k;

        // Initialize dp array with -1
        int[] dp = new int[1 << n];
        Arrays.fill(dp, -1);

        // Sum of empty subset is zero
        dp[0] = 0;

        // Iterate over all subsets/masks
        for (int mask = 0; mask < (1 << n); mask++) {
          
            // if current mask is invalid, continue
            if (dp[mask] == -1)
                continue;

            // Iterate over all array elements
            for (int i = 0; i < n; i++) {
              
                // Check if the current element can be added
                // to the current subset/mask
                if ((mask & (1 << i)) == 0
                    && dp[mask] + arr[i] <= target) {
                  
                    // transition
                    dp[mask | (1 << i)]
                        = (dp[mask] + arr[i]) % target;
                }
            }
        }

        // If the dp value of all elements used is zero,
        // then partitioning is possible
        return dp[(1 << n) - 1] == 0;
    }

    
    public static void main(String[] args) {
      
        int[] arr = { 2, 1, 4, 5, 3, 3 };
        int k = 3;

        if (isKPartitionPossible(arr, k)) {
            System.out.println("true");
        }
        else {
            System.out.println("false");
        }
    }
}
Python
# Program program to check if the
# given array can be partitioned
# into K subsets with equal sum
# using bitmasking dp


def isKPartitionPossible(arr, k):
    n = len(arr)

    if k == 1:
      
        # Return True as the entire array 
        # is the answer
        return True

    # If total number of partitions exceeds 
    # size of the array
    if n < k:
        return False

    totalSum = sum(arr)

    # If the array sum is not divisible by K
    if totalSum % k != 0:
      
        # No such partitions are possible
        return False

    # Required sum of each subset
    target = totalSum // k

    # Initialize dp array with -1
    dp = [-1] * (1 << n)

    # Sum of empty subset is zero
    dp[0] = 0

    # Iterate over all subsets/masks
    for mask in range(1 << n):
      
        # If the current mask is invalid, continue
        if dp[mask] == -1:
            continue

        # Iterate over all array elements
        for i in range(n):
          
            # Check if the current element can be added 
            # to the current subset/mask
            if (mask & (1 << i)) == 0 and dp[mask] + arr[i] <= target:
              
                # Transition
                dp[mask | (1 << i)] = (dp[mask] + arr[i]) % target

    # If the dp value of all elements used is zero,
    # then partitioning is possible
    return dp[(1 << n) - 1] == 0


if __name__ == "__main__":
    arr = [2, 1, 4, 5, 3, 3]
    k = 3

    if isKPartitionPossible(arr, k):
        print("true")
    else:
        print("false")
C#
// C# program to check if the
// given array can be partitioned
// into K subsets with equal sum
// using bitmasking dp

using System;

class GfG {
 
    static bool isKPartitionPossible(int[] arr, int k) {
        int n = arr.Length;

        if (k == 1) {

            // Return true as the entire array 
          	// is the answer
            return true;
        }

        // If total number of partitions exceeds
      	// size of the array
        if (n < k) {

            return false;
        }

        int totalSum = 0;
        foreach(int num in arr) { totalSum += num; }

        // If the array sum is not divisible by k
        if (totalSum % k != 0) {
            return false;
        }

        // Required sum of each subset
        int target = totalSum / k;

        // Initialize dp array with -1
        int[] dp = new int[1 << n];
        for (int i = 0; i < (1 << n); i++) {
            dp[i] = -1;
        }

        // Sum of empty subset is zero
        dp[0] = 0;

        // Iterate over all subsets/masks
        for (int mask = 0; mask < (1 << n); mask++) {

            // If the current mask is invalid, continue
            if (dp[mask] == -1)
                continue;

            // Iterate over all array elements
            for (int i = 0; i < n; i++) {

                // Check if the current element can be added
                // to the current subset/mask
                if ((mask & (1 << i)) == 0
                    && dp[mask] + arr[i] <= target) {

                    // Transition
                    dp[mask | (1 << i)]
                        = (dp[mask] + arr[i]) % target;
                }
            }
        }

        // If the dp value of all elements used is zero,
        // then partitioning is possible
        return dp[(1 << n) - 1] == 0;
    }

    static void Main(string[] args) {

        int[] arr = { 2, 1, 4, 5, 3, 3 };
        int k = 3;

        if (isKPartitionPossible(arr, k)) {
            Console.WriteLine("true");
        }
        else {
            Console.WriteLine("false");
        }
    }
}
JavaScript
// JavaScript program to check if the
// given array can be partitioned
// into K subsets with equal sum
// using bitmasking dp

function isKPartitionPossible(arr, k) {
    const n = arr.length;

    if (k === 1) {
    
        // Return true as the entire 
        // array is the answer
        return true;
    }

    // If total number of partitions exceeds size of the
    // array
    if (n < k) {
        return false;
    }

    let totalSum = 0;
    for (let num of arr) {
        totalSum += num;
    }

    // If the array sum is not divisible by k
    if (totalSum % k !== 0) {
        return false;
    }

    // Required sum of each subset
    const target = totalSum / k;

    // Initialize dp array with -1
    const dp = new Array(1 << n).fill(-1);

    // Sum of empty subset is zero
    dp[0] = 0;

    // Iterate over all subsets/masks
    for (let mask = 0; mask < (1 << n); mask++) {
    
        // If the current mask is invalid, continue
        if (dp[mask] === -1)
            continue;

        // Iterate over all array elements
        for (let i = 0; i < n; i++) {
        
            // Check if the current element can be added to
            // the current subset/mask
            if ((mask & (1 << i)) === 0
                && dp[mask] + arr[i] <= target) {
                
                // Transition
                dp[mask | (1 << i)]
                    = (dp[mask] + arr[i]) % target;
            }
        }
    }

    // If the dp value of all elements used is zero, then
    // partitioning is possible
    return dp[(1 << n) - 1] === 0;
}

const arr = [ 2, 1, 4, 5, 3, 3 ];
const k = 3;

if (isKPartitionPossible(arr, k)) {
    console.log("true");
}
else {
    console.log("false");
}

Output
true


Next Article

Similar Reads