Problem: Angry Children 2
You have to pick K elements from N elements, so that
is minimized.
Idea: Take all possible K elements and calculate result with the given equation. print the minimum one among them. (TLE!!)
Optimization 1: You don't need to take all possible K elements, look, your task is to minimize the result, so take consecutive K elements. that means you have to sort the array first. [if u have no idea about sliding window, read the previous post
here.
now you have selected all possible valid K elements in O(N)
Optimization 2: let, K=4 and the array after sorting is:-
[x0, x1, x2, x3, x4, x5, x6]
so, first segment is:-
[x0, x1, x2, x3], here, st=0, ed=3.
you can calculate result for this segment according to the given equation is:-
for(int i=st; i<ed; i++)
for(int j=i+1; j<=ed; j++)
res+=arr[j] - arr[i];
your can see, you should get TLE with such wired code even for small input.
lets see that above code actually does.
( x1 - x0 ) + ( x2 - x0 ) + ( x3 - x0 ) +
( x2 - x1 ) + ( x3 - x1 ) +
( x3 - x2 )
after simplify, we can get this
- 3*x0 - x1 + x2 + 3*x3
i think you already notice. if not, 1st coefficient is k*(-1)+1, then coefficient increase with +2 each time.
the simplified code will look like this:-
ml = k * (-1) + 1;
for( int i=st; i<=ed; i++, ml+=2)
res+= (arr[i] * ml);
Optimization 3: u thought u finished? nope... N<=10^5. you will still get TLE for larger input.
lets see what happens for 2nd segment.
- 3*x1 - x2 + x3 + 3*x4
now look at the difference with the formula of 1st segment and 2nd segment.
by subtracting formula of 1nd segment from 2st segment we get.
- 3*x1 - x2 + x3 + 3*x4
(-) - 3*x0 - x1 + x2 + 3*x3
-------------------------------------------------------
3*x0 - 2*x1 - 2*x2 - 2*x3 + 3*x4
==>> 3*x0 + 3*x4 - 2 * (x1 + x2 + x3)
we can pre-calculate the cumulative to have summation of (X
I + X
i+1 ….. + X
k-1) in O(1).
use this optimization for 2nd to last segment.
Note: answer can be a very large (long long) value.
Code:
- #include<bits/stdc++.h>
- using namespace std;
-
- int main()
- {
- int arr[100001], cs[100001], k, n;
- long long res=0, mn, ml;
- {
- cin>>n>>k;
- for(int i=0;i<n;i++)
- cin>>arr[i];
- sort(arr, arr+n);
-
- cs[0]=arr[0];
- for(int i=1;i<n;i++) cs[i]=cs[i-1]+arr[i];
-
- for(int i=0, t=i+k-1; t<n;i++, t++)
- {
- if(i)
- {
- res+=(arr[i-1]*(k-1)+arr[t]*(k-1));
- res-=((cs[t-1]-cs[i-1])*2);
- mn=min(mn, res);
- }
- else
- {
- ml=k*(-1)+1;
- for(int j=i;j<=t;j++, ml+=2)
- res+=((long long)arr[j] * ml);
- mn=res;
- }
- }
- cout<<mn<<endl;
- }
- return 0;
- }