[Tutorial] Snake Eating Problem of Codechef Snackdown 2017 Qualification round

Abridged problem description:

Given the length of $n$ snakes, $a_0, a_1, a_2, ......... a_{n-1}$.

A snake $i$ can eat another snake $j$, if $a_j \leq a_i$ and after eating a snake, it's length is increased by 1.

Given $q$ queries. Each query contains an integer $x$, you have to tell what is the maximum number of snakes you can get which are of length at least $x$.

Problem link: SNAKEEAT

Soultion: 

Let store the lengths in an array and sort the array in increasing order. Since now the array is sorted we can use binary search or lower bound to search any value in this array.

Before processing the queries, let first observe something which will be helpful in processing the queries.

Observe that the number of snakes whose length $\geq x$ can be added to the answer immediately.
For any the other snake of length $a_i$, it needs to eat $(x - a_i)$ snakes to make it's length $x$.

So, how many snakes are there whose length is $\geq x$ ? We can use lower bound to find this. Let this index be $k$

But what should we do to find the rest ?

We might iterate from the index  $k-1$ and check whether we can increase it's length to $x$ or not.
From index $k-1$, for each length $a_i$ it need to eat $(x - a_i)$ snakes to make it's length $x$. Of course snakes should be eaten from the left side(staring from index $0$).
We stop when there is not enough snakes on the left side.

But this process will get TLE. Because complexity of a single such process can be $O(q*(n+log n))$ in worst case.

So what to do ?

Their may be lot of ideas to optimize it. Here I am sharing my solution idea.

 I would like to discuss the solution with an example.
 Let n=13 and the given lengths(after sorting) are-

Length:	1	1	3	5	6	8	8	12	13	14	15	15	20
Index:	0	1	2	3	4	5	6	7	8	9	10	11	12

Observe that for each index $i$, there are $i$ snakes on it's left side.
Let we are querying for $x = 15$.
If we use lower_bound for 15, we get the index $k$ = 10. So there are $n-k = 13-10 = 3$ snakes whose length is at least $x$, so this value can be added to our answer immediately. So, our current answer is, $ans = 3$.

Now, we have to find how many snakes can be made of size 15 by eating other snakes.

We have to find how many snakes from index $k-1 = 9$, we can make of length 15 by eating other snakes.
Index = 9, length  = 14, need to eat = 15-14 = 1. total need till now  = 1 (possible, 9 snakes left)
Index = 8, length  = 13, need to eat = 15-13 = 2. total need till now  = 3 (possible, 8 snakes left)
Index = 7, length  = 12, need to eat = 15-12 = 3. total need till now  = 6 (possible, 7 snakes left)

Index = 6, length  = 8, need to eat = 15-8 = 7. total need till now  = 13, it is not possible, because there are 6 snakes on the left side of index 6.
So, 3 snakes can make their length 15.
So, final ans = 3 + 3 = 6.

We stopped at that index where sum of needed snakes to eat is less than the amount of snakes on the left side of that index. For each query we are finding this index by iterating linearly. So, in worst case it takes $O(n)$ time. So, for each query it takes $O(log n + n)$ time. Total query processing time $O(q*(log n + n))$. Sorting takes $O(n log n)$ time. So, total time complexity for each test case is $O(n log n + q*(log n + n))$. Since $1\leq n, q \leq 10^5$, it must get TLE.

So, the question is: how can we improve it anyway?

Observe that from the index $k - 1$, up to index $i$, sum of needed snakes is calculated as $(x - a_{k-1}) + (x - a_{k-2}) + (x - a_{k-3}) + ...... (x-a_i)$.
If we could know, somehow, for any index $i$, sum needed snakes to make all the snakes from index $k-1$ to $i$ of length $x$ in O(1), then we could use binary search(or lower_bound) the index up to which we can make their length equals to $x$.

How can we do it ? It may seems difficult, since for each query the value of $x$ is changed. But we still can do it by following the below technique.

We can store the cumulative sum of the difference between each length and and length $a_{n-1}$, i.e. the largest length.

According to our example, we will store the sum of difference between each length and $a_{n-1} = 20$.

The cumulative sum array will be:

Cumulative sum :	139	120	101	84	69	55	43	31	23	16	10	5	0
Index:	0	1	2	3	4	5	6	7	8	9	10	11	12

Since the cumulative sum array is increasing from right to left, binary search is applicable.

We can now calculate the sum from any index $i$ to any index $j$ ($i\geq j$), in O(1), by sum[i] - sum[j+1].

But how to handle, if we query for $x = 15$, or any other value of $x$?

We calculated the cumulative sum with respect to 20. Difference, $d$ between $a_{n-1} = 20$ and $x = 15$ is 5.

So, for each index we added extra $d=5$. From $i$ to $j$ ($i\geq j$) there are total $i-j+1$ elements. So for this range we added extra sum = $5*(i-j+1)$

So, for any query $x$, the difference $d = a_{n-1} - x$

So, the needed sum of snakes, with respect to $x$, from any index $i$ to any index $j$ ($i\geq j$) is (sum[j] - sum[i+1]) - d*(i-j+1)$.

That's great !! We can now calculate the needed sum from index $k-1$ up to any index $i$ in $O(1)$.

Now, we have to use binary search(or lower_bound) to find the index up to which the needed sum is greater than or equal to the number of elements of it's left side.

If that index is $i$ and our first lower_bound index is $k$, then our answer to the query will be $n-k$ (from lower_bound) $ + k-i$ (from the 2nd binary search (or lower_bound).

That's all !!

Now the complexity of each query is $O(log n + log n)$

So, total complexity for each test case is $O(n log n + q*(log n + log n))$ = $O(n log n + 2*q* log n )$

Finally, here is the code(in c++):

/*=================================*\
                                   
      Md. Shahidul Islam           
         CSE, BRUR                 
      Rangpur, Bangladesh          
 mail: shahidul.cse.brur@gmail.com 
 FB  : fb.com/shahidul.brur        
 Blog: shahidul-brur.blogspot.com 
\*=================================*/
#include<bits/stdc++.h>
using namespace std;

#define ll long long

const int MX = 100005;
int a[MX], n;
ll sum[MX];

int cal(int x, int last)
{
    if(last==0)
        return 0;
    int diff = a[n-1] - x;
    int l = 0, r = last;
    int ans = 0;
    while(l<=r)
    {
        int mid = (l+r)/2;
        if(mid==0)
        {
            l = mid+1;
            continue;
        }
        
        ll need = (sum[last] - sum[mid-1]) - 1LL*diff*(last - mid + 1);
        if(mid>=need)
        {
            ans = last-mid+1;
            r = mid-1;
        }
        else l = mid+1;
    }
    return ans;
}
int main()
{
    //ios_base::sync_with_stdio(false); cin.tie(NULL);
    int t, q, x;
    cin>>t;
    while(t--)
    {
        cin>>n>>q;
        for(int i = 0; i< n; i++)
        {
            cin>>a[i];
        }
        sort(a, a+n);
        sum[0] = 0LL;
        for(int i = 1; i< n; i++)
            sum[i] = 1LL*sum[i-1] + 1LL*(a[n-1] - a[i]);
        
        while(q--)
        {
            cin>>x;
            int cnt = 0;
            int idx = lower_bound(a, a+n, x) - a;
            cnt=(n-idx);
            cnt+=cal(x, idx-1);
            cout << cnt << "\n";
        }
    }
    return 0;
}