Split an Array A[] into Subsets having equal Sum and sizes equal to elements of Array B[]

Given an array A[] consisting of N integers, the task is to split the array A[] into subsets having equal sum and of length equal to elements in array B[].


Input: A[] = {17, 13, 21, 20, 50, 29}, B[] = {2, 3, 1}
21 29
17 13 20

Input: A[] = { 1, 2, 3, 4, 5, 6}, B[] = { 2, 2, 2}
1 6
2 5
3 4

Approach: Follow the steps below to solve the problem:

  • Since the count of subsets is already given, calculate the sum of each subset.
  • Traverse through each element of array B[], find each possible combination of length B[i] and check if the sum of the combination is equal to the desired sum or not.
  • Repeat the above steps for every array element B[i] and print the final answer

Below is the implementation of the above approach:






# Python Program to implement
# the above approach
from itertools import combinations
# Function to split elements of first
# array into subsets of size given as
# elements in the second array
def main(A, B):
    # Required sum of subsets
    req_sum = sum(A) // len(B)
    # Stores the subsets
    final = []
    # Iterate the array B[]
    for i in B:
        # Generate all possible
        # combination of length B[i]
        temp = list(combinations(A, i))
    # Iterate over all the combinations
        for j in temp:
            # If the sum is equal to the 
            # required sum of subsets
            if(sum(j) == req_sum):
                # Store the subset
                temp = list(j)
                for k in temp:
                    # Removing the subset
                    # from the array
    # Printing the final result
# Driver Code
if __name__ == '__main__':
    # Value of array A and B
    A = [1, 2, 3, 4, 5, 6]
    B = [2, 2, 2]
    # Function call
    main(A, B)



[[1, 6], [2, 5], [3, 4]]

Time Complexity: O(N3)
Auxiliary Space: O(2maxm), where maxm is the maximum element in array B[]

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 Posts:

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.