Home » Interview coding problems/challenges

Partition to K Equal Sum Subsets

Here, we are going to learn about the solution of partition to k equal sum subsets and its C++ implementation.
Submitted by Divyansh Jaipuriyar, on August 16, 2020

Problem statement

Given an array of integers A[] and a positive integer k, find whether it's possible to divide this array into k non-empty subsets whose sums are all equal.

The problem wants you to find if it is possible to partition given array into k subset such that each set has sum equal to (total sum of array/number of the required partition), keeping in mind that we can use one element in only one of the subset. If it is possible we are required to print "YES" otherwise "NO".

Input

The first line of input consists of T number of test cases, following each test case, consist of N the size of the array following lined consist of N elements of the array.

Output

You need to determine if it possible to divide the array into K subset of equal sum or not.

Examples

Input:
T = 1
N = 5, K = 3
[2,1,4,5,6]

Output:
YES, as {1,5},{4,2} and {6}

Solution Approach

Since we have the probability of either including the element into the subset or not including it which means we can choose some other element from the array for current element and place the other elements into some other subset, this implies the concept of backtracking.

So we will use the following steps:

1. First, we would check if the sum of the elements of the array is divisible by k or not
2. If it is not divisible by k then simply return false otherwise proceed further, we will add each number of the array A[] into one of k group as long as the sum of that group is not larger than the (sum/k) value.
3. For each of these process, we recursively search with one less number to consider in A[], if we are able to add all elements of the A[] then we return true, otherwise, we return false.
4. In order to make the process faster, we will sort the A[] array so that we first add the largest elements first.
5. After sorting the array we will traverse from the last element to the first element order. We will create a boolean function find which will return true or false depending upon valid partition. The parameters used in the function are given array, the vector in which we are adding elements, the target value, and the index that we are using from the last.

If our index reached <0 then it means that all elements of the given array have been used in some of the subset for partition and hence we return true.

Otherwise, we used current value and add it to vector and make a recursive call to check if leads to a solution or not, if not then we subtract that value and backtrack, if that index position in vector(use for partition) is still 0 then it means that partition is not possible and we return false.

Time Complexity for the above approach is: O(k^(n-k)K!), n is array size and k is the required partition size.

Space Complexity for the above approach is : O(n)

C++ Implementation

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

typedef long long ll;

//declare find function.
bool find(vector<ll>& p, ll arr[], ll r, ll tar)
{
    //if we successfully used all array.
    //elements to construct array then return true.
    if (r < 0)
        return true;

    //use val variable for current index value.
    ll val = arr[r];
    //decrease last array elements.
    r--;

    //iterate through the p vector.
    for (ll i = 0; i < p.size(); i++) {

        //check for current value val.
        if (val + p[i] <= tar) {
            //add it to vector.
            p[i] += val;
            //call find function.
            if (find(p, arr, r, tar))
                return true;
            //backtrack if not possible from current val.
            p[i] -= val;
        }

        //if current subset is 0 then break.
        if (p[i] == 0)
            break;
    }
    //return false if not possible.
    return false;
}

int main()
{
    ll t;

    cout << "Enter number of test cases: ";
    cin >> t;

    while (t--) {
        cout << "Enter N and K: ";

        ll n, k;
        cin >> n >> k;

        ll arr[n];

        cout << "Enter array elements: ";
        for (ll i = 0; i < n; i++)
            cin >> arr[i];
        //take sum of elements.
        ll sum = accumulate(arr, arr + n, 0);

        //check if sum%k is 0 or not
        //if it is not divisible then return false.
        if (sum % k) {
            cout << "NO"
                 << "\n";
        }
        else {
            //store required target value in tar.
            ll tar = sum / k;

            //sort the array.
            sort(arr, arr + n);

            //initialize index from last pos.
            ll r = n - 1;

            //check greatest value is greater than required than
            //target value then return false.
            if (arr[r] > tar) {
                cout << "NO"
                     << "\n";
                continue;
            }

            //keep decreasing indexes if last value
            //is equal to target value.
            while (arr[r] == tar) {
                r--;
                k--;
            }
            //initialize vector for subset.
            vector<ll> p;
            //fill all with 0 values.
            for (ll i = 0; i < k; i++)
                p.push_back(0);
            if (find(p, arr, r, tar))
                cout << "YES"
                     << "\n";
            else
                cout << "NO"
                     << "\n";
        }
    }
    return 0;
}

Output

Enter number of test cases: 3
Enter N and K: 5 3
Enter array elements: 2 1 4 5 6
YES
Enter N and K: 6 2
Enter array elements: 1 3 4 8 6 2
YES
Enter N and K: 5 2
Enter array elements: 1 3 5 7 9
NO