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[].

Examples:

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

Input: A[] = { 1, 2, 3, 4, 5, 6}, B[] = { 2, 2, 2}
Output:
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

# 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)
                final.append(temp)

                for k in temp:

                    # Removing the subset
                    # from the array
                    A.remove(k)
                break

    # Printing the final result
    print(final)


# 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)

Output:

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

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