Editorial for COCI '10 Contest 1 #4 Ljutnja
Submitting an official solution before solving the problem yourself is a bannable offence.
Let be the sum of all wishes. Imagine that each child was given as many candies as it wants, and now they have to be taken away, leaving candies in total. That is, we have to take candies away.
Since the sum of the numbers of candies that the children will be missing is constant (equal to ), and we wish to minimize the sum of squares of those numbers, a key observation is that the numbers should be "as equal as possible". More precisely, it is possible to prove mathematically (using the inequality of arithmetic and geometric means) that the sum of squares of positive numbers, with a fixed sum, is minimal when they are equal.
Let be the remaining number of candies that need to be taken away (the starting value of is ). Our goal is to, if possible, take away an equal number of candies from each child, specifically . If this is not possible, the number of candies taken away should be as close as possible to .
If the child that wants the smallest amount of candies has at least candies, we will be able to fulfill the goal by taking candies away from it, thus reducing the problem to the remaining children (by decrementing , calculating the new value of , and processing the child with the next smallest wish). On the other hand, if the observed child has less than candies, preventing us from taking that many from it, we will take as many as we can, i.e. all the candies it was given, and proceed to process the remaining children.
In the end, we can just output the sum of squared numbers of candies taken away from the children. The complexity of the described algorithm is linear, but it requires first sorting children's wishes into a nondecreasing sequence, resulting in the total complexity of .
It is also possible to solve this problem using binary search, however this solution is left as an exercise for the reader.
Comments