Dividing array into chunks of almost equal sum

We can first try to divide an array into 2 parts such that their sum is almost equal.
Then once we have the 2 sets, we can further apply the same process on each of them to get 2*2 = 4 sets of equal sum.

The algorithm to divide an array into 2 parts of approximately equal sum is as follows:

• Initialize 2 empty sets to hold our answer.
• Sort the array in reverse order.
• While maintaining the sum of the 2 sets, iterate over all the elements from array and append them into the set which has the lesser sum.
• Note that this is just an approximate algorithm. If we want to find the exactly optimal answer, then we can model this problem as the subset sum problem, and find if we can divide the array into 2 parts, where one of the set has the sum sum/2 or sum/2 - 1 or sum/2 - 2 ... 0 (trying each of them in that order). This is significantly slower compared to our approximate solution.

def divide_almost_equally_into_2(arr):
    set1 = []
    set2 = []
    sum1 = sum2 = arr_idx = 0
    while arr_idx < len(arr):
        if sum1 < sum2:
            set1.append(arr[arr_idx])
            sum1 += arr[arr_idx]
        else:
            set2.append(arr[arr_idx])
            sum2 += arr[arr_idx]
        arr_idx += 1
    return set1, set2


def divide_almost_equally_into_4(arr):
    arr.sort(reverse=True)
    set1, set2 = divide_almost_equally_into_2(arr)
    set11, set12 = divide_almost_equally_into_2(set1)
    set21, set22 = divide_almost_equally_into_2(set2)
    return [set11, set12, set21, set22]


def main():
    arr = [11,20,2,4,8,13,16,0,1,0,3,6]
    set1, set2, set3, set4 = divide_almost_equally_into_4(arr)
    print(f"{arr}   {sum(arr)}\n")
    print(f"{set1}   {sum(set1)}")
    print(f"{set2}   {sum(set2)}")
    print(f"{set3}   {sum(set3)}")
    print(f"{set4}   {sum(set4)}")


main()

Output:

[20, 16, 13, 11, 8, 6, 4, 3, 2, 1, 0, 0]   84

[13, 8]   21
[16, 3, 2]   21
[11, 6, 4]   21
[20, 1, 0, 0]   21

EDIT:

To generalize the same algorithm to 'n' number of splits, we can use a heap:

• Create a min-heap of size 'n', where each element is a tuple of the form (current_sum_of_set_i, i).
• So, initially heap will contain elements (0, 0), (0, 1) ... (0, n-1).
• Now iterate over the reverse-sorted array, and assign each element to the set present at the top of the heap.
• Update the heap element of the set with the new sum of the element added to it.

import heapq


def divide_almost_equally(arr, num_chunks):
    arr = sorted(arr, reverse=True)
    heap = [(0, idx) for idx in range(num_chunks)]
    heapq.heapify(heap)
    sets = {}
    for i in range(num_chunks):
        sets[i] = []
    arr_idx = 0
    while arr_idx < len(arr):
        set_sum, set_idx = heapq.heappop(heap)
        sets[set_idx].append(arr[arr_idx])
        set_sum += arr[arr_idx]
        heapq.heappush(heap, (set_sum, set_idx))
        arr_idx += 1
    return sets.values()


def main():
    arr = [11,20,2,4,8,13,16,0,1,0,3,6]
    set1, set2, set3, set4 = divide_almost_equally(arr, 4)
    print(f"{sorted(arr, reverse=True)}   {sum(arr)}\n")
    print(f"{set1}   {sum(set1)}")
    print(f"{set2}   {sum(set2)}")
    print(f"{set3}   {sum(set3)}")
    print(f"{set4}   {sum(set4)}")


main()

Output:

[20, 16, 13, 11, 8, 6, 4, 3, 2, 1, 0, 0]   84

[20, 1]   21
[16, 4, 0, 0]   20
[13, 6, 3]   22
[11, 8, 2]   21

Here is a solution where the sum of each chunk is less than or equal to a threshold. This solution is inspired by the algorithm proposed by Anmol Singh Jaggi.

In the function below, the input set is divided recursively till the summation of each chunk is less than or equal to the Threshold

The number of the chunk is increasing linearly, but you could change according to your preference, such as num_chunks = num_chunks*2

def divide_almost_equally(arr, num_chunks,Threshold):
    divideMore = 0
    arr = sorted(arr, reverse=True)
    set_sum = np.zeros(num_chunks)
    sets = {}
    for i in range(num_chunks):
        sets[i] = []
    arr_idx = 0
    while arr_idx < len(arr):
        ChunkIndx = np.argmin(set_sum)
        sets[ChunkIndx].append(arr[arr_idx])
        set_sum[ChunkIndx] = set_sum[ChunkIndx] + arr[arr_idx]
        arr_idx = arr_idx + 1
    for total in set_sum:
        if(total>Threshold):
            divideMore = 1
            break
    if(divideMore == 1):
      num_chunks = num_chunks + 1
      return divide_almost_equally(arr, num_chunks,Threshold)
    else:
        return sets, num_chunks

if __name__ == "__main__":
   input = np.arange(0,100)
   Set, num_chunks = divide_almost_equally(np.arange(0,100), 1,100) # start with num_chunks = 1num_chunks = 1
   for chunk_indx in range(num_chunks): 
      print(Set[chunk_indx],'-- >',np.sum(Set[chunk_indx]))