Best solutions for Microsoft interview tasks. Partition array into N subsets with balanced sum.

Description:

Give you one sorted array, please put them into n buckets, we need to ensure we get n sub array with approximately equal weights.

Example:

input {1, 2, 3, 4, 5} n = 3

output [[5],[1,4],[2,3]];

Solution:

Generally this task look like a typical partition problem https://en.wikipedia.org/wiki/Partition_problem

Unfortunately this description looks unclear or too general because it doesn’t describe what exactly “approximately equal weights” means, so I don’t believe this question will be asked on an interview this way. You will need to ask the interviewer about details and choose the right approach.

There are many approaches to solve this problem optimized for special cases. In case we know nothing special about the input data and the requirements, so in my opinion the best solution is to use the simplest and the most basic approach — the greedy algorithm. https://en.wikipedia.org/wiki/Greedy_algorithm

The greedy algorithm is based on a simple idea: we take the biggest value from the input array and put it to the partition with the lowest sum. Our input array is already sorted by description so we may take the vales one by one. In order to pick a partition with lowest sum we may keep the partitions in priority queue which is sorted by sums of these partitions.

C++ code:

vector<vector<int>> part(const vector<int> & v, int n) {
    int v_size = v.size();
    // in this vector we will keep sums of the partitions
    vector<int> sums(n, 0);
    
    // in this priority queue we will keep values of the priorities
    // it is just numbers from 0 to numbers of partitions. 
    // We will use these numbers as indexes  
    // to access sums and to definition priorities of the partitions
    // in other words the sums tied with the partitions by 
    // these priority values.
    
    // Pay attention that the sums we keep in the external 
    // vector described above, 
    // so we need to declare a special comparator
    // for the queue in order to sort the queue by the 
    // priority values.
    auto compare = [&sums](int a, int b) { 
        return sums[a] >= sums[b]; 
    };
    priority_queue<int, std::vector<int>, 
                   decltype(compare)> pq(compare);
    // here we will keep values of the partitions
    vector<vector<int>> result(n, vector<int>());
    // generate numbers of the priority values.
    for (int i = 0; i < n; i++) {
        pq.push(i);
    }
    // The input array is sorted ascending so we will process 
    // it from the back to the front.
    for (int i = v_size - 1; i >= 0; --i) {
        // take the highest priority from the queue
        int c = pq.top(); pq.pop();
        // put the max value from the input array to the partition 
        // with the highest priority, 
        // that is partition with the lowest sum.
        result[c].push_back(v[i]);
        // increase the lowest sum to max value from 
        // the input array.    
        sums[c] += v[i];
        // push back the priority value. After this call the queue 
        // will be resorted corresponding to 
        // new value of the sums array.
        pq.push(c);
    }
    return result;
}

Repository with the full project you can find here: https://github.com/jolly-fellow/microsoft/tree/master/partition_array_into_n_subsets_with_balanced_sum

Return to the table of contents.