Approach
To solve this problem, we need to transform this problem into a graphical one. Let's first sort the array in nonincreasing order. Now, let's assume each array element represents a node in a kary tree. A kary tree is a tree where each node has at most ‘K’ children.
We can visualise the operation of adding an array element, ‘ELEM1’ to another element ‘ELEM2’ as appending the subtree of ‘ELEM1’ as a child to ‘ELEM2’. We can append the subtree of ‘ELEM1’ to ‘ELEM2’ only if ‘ELEM2’ has fewer than ‘K’ children.
Now, as we have identified an operation in trees that resemble the addition operation described in the problem, let's move on to the definition of cost in the new context.
The cost of collecting the sum of an array can be calculated by summing up the multiplication of each node with its depth in the tree, assuming the root is at depth 0.
To understand this, let's take the example mentioned before.
Let's create the tree corresponding to the operations. Each previously mentioned operation translates to an edge as shown in the images below.
The initial configuration
Operation 1
Operation 2
Operation 3
Operation 4
Cost = 8 * 0 + 4 * 1 + 6 * 1 + 2 * 2 + 3 * 2 = 20
Now, the question is how to minimize the cost?
Taking intuition from the tree structure, we can do the following:
 Nodes with higher values should have lower depth.
 Each nonleaf node should have K children to accumulate as many nodes as possible in the lower depths.
Remember, we need not create the tree. We just did this analysis to understand the problem better.
Algorithm
 Take the input array and integer K.
 Sort the array in nonincreasing order.
 Calculate the prefix sum in the ‘PREF_SUM’ array.

Initialise the following variables:
 COST: To store the cost of collection.
 START_INDEX: It will depict the index of the leftmost node at each level.
 DEPTH: To store the current depth as the algorithm runs.
 DEPTH_SIZE: To hold the number of nodes at each depth.

While the ‘START_INDEX’ is less than ‘N’, do the following:
 Calculate END_INDEX = START_INDEX + DEPTH_SIZE.
 Take the minimum of ‘END_INDEX’ and ‘N  1’ and store it in ‘END_INDEX’
 Calculate the sum of elements in the range described by ‘START_INDEX’ and ‘END_INDEX’ using ‘PREF_SUM’ array.
 Increment ‘COST’ by the sum calculated in previous step multiplied by the current depth.
 Increment ‘DEPTH’ by one.
 DEPTH_SIZE *= K to take into account the second point on minimising the cost.
 Output the COST.
Program
#include<iostream>
#include<algorithm>
using namespace std;
int findCost(int arr[])
{
// Size of the array.
int N = sizeof(arr)/sizeof(arr[0]);
// Integer K as described in the problem statement.
int K = 2;
// Sort the array in the nonincreasing order.
sort(arr, arr + N, greater<int>());
// Prefix sum calculation.
int prefSum[N+1];
prefSum[0] = 0;
// Standard technique to compute prefix sum in O(N) time.
for(int i = 0; i < N; i++)
{
prefSum[i + 1] = prefSum[i] + arr[i];
}
// Cost of sum collection.
int cost = 0;
// Elements to pick from. 0th index is ignored as it will always be multiplied by zero, thus contribution to the cost.
int startIndex = 1;
// Current depth.
int depth = 1;
// Number of nodes at the current depth.
int depthSize = K;
while(startIndex < N)
{
// Array elements covered at this depth.
int endIndex = startIndex + depthSize  1;
// In case of overflow.
endIndex = min(endIndex, N  1);
// Sum of elements at the current depth.
int levelSum = prefSum[endIndex + 1]  prefSum[startIndex];
// As described by the approach.
cost += levelSum * depth;
// Update the startIndex after covering all the elements till endIndex.
startIndex = endIndex + 1;
// The number of nodes at the next level will be K times the number of nodes at this level.
// This is to accumulate as many nodes as possible in the lower depths.
// Level is used interchangeably with depth.
depthSize *= K;
// Increase the depth.
depth++;
}
}
int main()
{
// Read the size.
int N;
cout << "Enter the size of the array: ";
cin >> N;
// Input array.
int arr[N];
cout << "Enter the elements: ";
for(int i = 0; i < N; i++){
cin >> arr[i];
}
// Output the answer.
cout << findCost(arr) << endl;
}
Input
Enter the size of the array: 5
Enter the elements: 4 8 6 2 3
Output
20
Time Complexity
The time complexity of the above approach is O(N * logN), where N is the size of the input array.
It is because we are sorting the array in nonincreasing order which takes O(N * log N). The time complexity of the while loop is O(log_{K}N) where K is the input parameter, however, the sorting algorithm is the upper bound to the time complexity.
Space Complexity
The space complexity of the above approach is O(N) to store the prefix sum array.
Check out this array problem  Merge 2 Sorted Arrays
Key Takeaways
This blog discussed a very interesting problem based on prefix sums and kary trees. It was intriguing to analyse the problem graphically. We created operations in the graphical domain that resemble operations described in the problem. The most important part and the motive of the construction was to understand how to minimise the cost of collecting the sum.
Hence learning never stops, and there is a lot more to learn.
So head over to our practice platform Coding Ninjas Studio to practice top problems, attempt mock tests, read interview experiences, and much more. Till then, Happy Coding!