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)

Java

import java.util.*;
import java.math.*;
 
public class Main {
 
    // Function to split elements of first
    // array into subsets of size given as
    // elements in the second array
    static void main(int[] A, int[] B) {
 
        // Required sum of subsets
        int req_sum = Arrays.stream(A).sum() / B.length;
 
        // Stores the subsets
        List<List<Integer>> finalList = new ArrayList<>();
 
        // Iterate the array B[]
        for (int i : B) {
 
            // Generate all possible
            // combination of length B[i]
            List<List<Integer>> tempList = subsets(A, i);
 
            // Iterate over all the combinations
            for (List<Integer> j : tempList) {
 
                // If the sum is equal to the
                // required sum of subsets
                if (j.stream().mapToInt(Integer::intValue).sum() == req_sum) {
 
                    // Store the subset
                    finalList.add(j);
 
                    // Removing the subset
                    // from the array
                    A = removeElements(A, j);
                    break;
                }
            }
        }
 
        // Printing the final result
        System.out.println(finalList);
    }
 
    private static int[] removeElements(int[] A, List<Integer> j) {
        List<Integer> result = new ArrayList<>();
        for (int i : A) {
            if (!j.contains(i)) {
                result.add(i);
            }
        }
        return result.stream().mapToInt(Integer::intValue).toArray();
    }
 
    private static List<List<Integer>> subsets(int[] A, int i) {
        List<List<Integer>> result = new ArrayList<>();
        combinations(A, 0, i, new ArrayList<>(), result);
        return result;
    }
 
    private static void combinations(int[] A, int start, int k, List<Integer> current, List<List<Integer>> result) {
        if (k == 0) {
            result.add(new ArrayList<>(current));
            return;
        }
        for (int i = start; i < A.length; i++) {
            current.add(A[i]);
            combinations(A, i + 1, k - 1, current, result);
            current.remove(current.size() - 1);
        }
    }
 
    // Driver Code
    public static void main(String[] args) {
        // Value of array A and B
        int[] A = new int[]{1, 2, 3, 4, 5, 6};
        int[] B = new int[]{2, 2, 2};
 
        // Function call
        main(A, B);
    }
}

Javascript

// Function to split elements of first
// array into subsets of size given as
// elements in the second array
function main(A, B) {
 
  // Required sum of subsets
  const req_sum = Math.floor(A.reduce((a, b) => a + b, 0) / B.length);
 
  // Stores the subsets
  const final = [];
 
  // Iterate the array B[]
  for (const i of B) {
 
    // Generate all possible
    // combination of length B[i]
    const temp = combinations(A, i);
 
    // Iterate over all the combinations
    for (const j of temp) {
 
      // If the sum is equal to the
      // required sum of subsets
      if(j.reduce((a, b) => a + b, 0) == req_sum){
 
        // Store the subset
        const subset = [...j];
        final.push(subset);
 
        for (const k of subset) {
 
          // Removing the subset
          // from the array
          A.splice(A.indexOf(k), 1);
        }
        break;
      }
    }
  }
 
  // Printing the final result
  console.log(final);
}
 
// Helper function to generate all possible
// combinations of length k from an array A
function combinations(A, k) {
  if (k === 0) {
    return [[]];
  }
  if (A.length === 0) {
    return [];
  }
  const [first, ...rest] = A;
  const combsWithoutFirst = combinations(rest, k);
  const combsWithFirst = combinations(rest, k - 1);
  const combsWithFirstMapped = combsWithFirst.map(comb => [first, ...comb]);
  return [...combsWithoutFirst, ...combsWithFirstMapped];
}
 
// Driver Code
const A = [1, 2, 3, 4, 5, 6];
const B = [2, 2, 2];
 
// Function call
main(A, B);

C#

using System;
using System.Collections.Generic;
using System.Linq;
 
class Program
{
    // Function to split elements of first
    // array into subsets of size given as
    // elements in the second array
    static void Main(int[] A, int[] B)
    {
        // Required sum of subsets
        int req_sum = A.Sum() / B.Length;
 
        // Stores the subsets
        List<List<int>> final = new List<List<int>>();
 
        // Iterate the array B[]
        foreach (int i in B)
        {
            // Generate all possible
            // combination of length B[i]
            List<int[]> temp = GetCombinations(A, i);
 
            // Iterate over all the combinations
            foreach (int[] j in temp)
            {
                // If the sum is equal to the
                // required sum of subsets
                if (j.Sum() == req_sum)
                {
                    // Store the subset
                    List<int> subset = j.ToList();
                    final.Add(subset);
 
                    foreach (int k in subset)
                    {
                        // Removing the subset
                        // from the array
                        A = A.Where(val => val != k).ToArray();
                    }
                    break;
                }
            }
        }
 
        // Printing the final result
        Console.WriteLine("[{0}]", string.Join(", ", final.Select(x => "[" + string.Join(", ", x) + "]")));
    }
 
    // Function to generate all possible
    // combinations of given length
    static List<int[]> GetCombinations(int[] arr, int len)
    {
        List<int[]> result = new List<int[]>();
 
        // Edge case
        if (len > arr.Length)
        {
            return result;
        }
 
        // Generate all possible combinations
        for (int i = 0; i < arr.Length; i++)
        {
            int[] temp = new int[len];
            temp[0] = arr[i];
 
            if (len == 1)
            {
                result.Add(temp);
            }
            else
            {
                List<int[]> nextCombos = GetCombinations(arr.Skip(i + 1).ToArray(), len - 1);
                foreach (int[] combo in nextCombos)
                {
                    combo.CopyTo(temp, 1);
                    result.Add(temp);
                    temp = new int[len];
                    temp[0] = arr[i];
                }
            }
        }
 
        return result;
    }
 
    // Driver code
    static void Main()
    {
        // Value of array A and B
        int[] A = new int[] { 1, 2, 3, 4, 5, 6 };
        int[] B = new int[] { 2, 2, 2 };
 
        // Function call
        Main(A, B);
    }
}
// This code is contributed by divyansh2212

C++

#include <bits/stdc++.h>
#include <vector>
#include <numeric> // for std::accumulate
 
using namespace std;
 
void combinations(vector<int>& A, int start, int k, vector<int> current, vector<vector<int>>& result) {
    if (k == 0) {
        result.push_back(current);
        return;
    }
    for (int i = start; i < A.size(); i++) {
        current.push_back(A[i]);
        combinations(A, i + 1, k - 1, current, result);
        current.pop_back();
    }
}
 
vector<vector<int>> subsets(vector<int>& A, int k) {
    vector<vector<int>> result;
    combinations(A, 0, k, {}, result);
    return result;
}
 
vector<int> removeElements(vector<int>& A, vector<int>& j) {
    vector<int> result;
    for (int i : A) {
        if (find(j.begin(), j.end(), i) == j.end()) {
            result.push_back(i);
        }
    }
    return result;
}
 
// Function to split elements of first
// array into subsets of size given as
// elements in the second array
void mainFunc(vector<int>& A, vector<int>& B) {
 
    // Required sum of subsets
    int req_sum = accumulate(A.begin(), A.end(), 0) / B.size();
 
    // Stores the subsets
    vector<vector<int>> finalList;
 
    // Iterate the array B[]
    for (int i : B) {
 
        // Generate all possible
        // combination of length B[i]
        vector<vector<int>> tempList = subsets(A, i);
 
        // Iterate over all the combinations
        for (vector<int> j : tempList) {
 
            // If the sum is equal to the
            // required sum of subsets
            if (accumulate(j.begin(), j.end(), 0) == req_sum) {
 
                // Store the subset
                finalList.push_back(j);
 
                // Removing the subset
                // from the array
                A = removeElements(A, j);
                break;
            }
        }
    }
 
    // Printing the final result
    cout<<"[";
    int i=0;
    for (vector<int> subset : finalList) {
            cout <<"["<<subset[0] << ","<<subset[1]<<"]";
            if(i<finalList.size()-1)
                cout<<", ";
            i++;
    }
            cout << "]";
 
}
 
// Driver Code
int main() {
    // Value of array A and B
    vector<int> A{1, 2, 3, 4, 5, 6};
    vector<int> B{2, 2, 2};
 
    // Function call
    mainFunc(A, B);
    return 0;
}