Given a sequence of N numbers (a1, a2, ..., an) you have to find the minimum number of operations to get a consecutive group of size K such that all the elements in this group have the same height. To achieve this you can increase or decrease the height of any element by 1, any number of times, each modification costs 1 operation.
Constraints: 5 <= N <= 10^6 3 <= K <= N 0 <= ai <= 10^6
Example: N = 6 K = 4 10 4 5 2 5 7
Here the solution is a group from position 2 (1-based) to position 5 (1-based), which costs 4 operations.
So far I have a solution in O(Nlog^2(N)) using a binary search + segment tree, but this timed out. Could you help me to improve this?
I suspect the value to which you equalize the k elements is the median value. Thats because as you decrese the target T it is going to be good to keep decreasing while there are more smaller elements than larger ones.
So for a window of size k you just keep an augmented BST like an AVL tree and you can find the median and the sum of the elements smaller / larger than that in log n time.
Hope it helps.
Best, Mircea
Oh I see, but I have never done something with that DS, could you give some guidance or some code? Thanks.
Try Peter Brass Advanced Data Structures for info on how to implement an AVL or a Red Black tree. It is a useful book.
You can do it only with the segment tree (that counts the number of non void elements and their sum). Add the first k numbers. Using the segment tree that counts the number of elements you can find in logn the position of the median. Then you only need two more sums to calculate the costs. After that you move the sliding window to the right. Nlogn final complexity.
On how to calculate the median: start at the root, lets say you are looking for the kth element. Look at the left children. If there are more or equal to k then your element is there. If there are less then k then it's on the right. (If its on the right you need to search for the k — x element, where x is how many there are on right.)
Another solution is to constantly have the median using 2 heaps. One that is a max heap and has the k/2 smallest elements and a min heap havin k/2 biggest elements. (You can figure out how to insert / remove elements)
Thanks for your help, both solutions are quite good.