In the qualification round of Codechef Snackdown 2017, I was able to solve all the problems and I had a great fun to solve the problem SNAKEEAT.

I was so much excited that I had given a post to express my feeling after solving this problem.

After that some people asked me on codeforces or facebook, how I had solved the problem. So, I had decided to write a post regarding my solution idea of this problem.

Before staring the post, I want to say you that I am not so great programmer, so I may make mistake in my post. And, this is my first tutorial post on codeforces, so forgive me if you find anything wrong. Also, English is not my native language, so I may make some grammatical mistake in my post. So, I beg your forgiveness in advance.

Abridged problem description:

Given the length of n snakes, a₀, a₁, a₂, .........a_n - 1. A snake i can eat another snake j, if a_j ≤ 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.

Soultion:

Prerequisites:

1. binary search
2. STL lower_bound
3. Cumulative sum or prefix sum or range sum

In my post, I won't discuss about these. I assume that you already know these, if not please learn them first.

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  ≥ 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  ≥ 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 + logn)) 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-

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(logn + n) time. Total query processing time O(q * (logn + n)). Sorting takes O(nlogn) time. So, total time complexity for each test case is O(nlogn + q * (logn + n)). Since 1 ≤ n, q ≤ 10⁵, 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:

We can now calculate the sum from any index i to any index j (i ≤ j), in O(1), by sum[j] — sum[i-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 ≤ j) there are total j - i + 1 elements. So for this range we added extra sum = 5 * (j - i + 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 ≤ j) is (sum[j] - sum[i - 1]) - d * (i - j + 1).

That's great !! We can now calculate the needed range sum from index k - 1 up to any index i in O(1) by sum[k-1] — sum[i-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 at 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(logn + logn) So, total complexity for each test case is O(nlogn + q * (logn + logn)) = O(nlogn + 2 * q * logn)

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

#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;
}

You also can read this post on my blog

If you find anything wrong in my post, please let me know. If you face any difficulty to understand any part of my post, feel free to ask me in the comment below. Also, if you have another idea to solve this problem, you can share that in the comment.

Thanks to all.