Lintcode 617: Maximum Average Subarray II

https://www.lintcode.com/problem/maximum-average-subarray-ii/note/171180

Description

中文English

Given an array with positive and negative numbers, find the maximum average subarray which length should be greater or equal to given length k.

It's guaranteed that the size of the array is greater or equal to k.

Have you met this question in a real interview?  

Example

Example 1:

Input:
[1,12,-5,-6,50,3]
3
Output:
15.667
Explanation:
 (-6 + 50 + 3) / 3 = 15.667

Example 2:

Input:
[5]
1
Output:
5.000

Notes:

This question is not an easy question, but greedy can solve it.

One way to think about this question is that: let's say we are now at the ith element position. If we put the ith element into consideration, we will have i more subarrays ending with the ith element. But we only care about the ones with a length no less than k.

So just need to record the previous sum results ending with the (i-1)th element and the corresponding length (since we need the average of them). Then update these results in both sum and length: sum will be added by nums[i], and length will be added by 1. And we will have a new subarray with size of k ending at the ith element. So we also need to take this into account.

One improvement is that, we can maintain a decreasing array in the sum. Since that the subarrays with the sum is smaller and the length is longer cannot give a larger average. So them can be ruled out directly.

During the scanning, we update the maximum average.

See the code below:

class Solution {
public:
    /**
     * @param nums: an array with positive and negative numbers
     * @param k: an integer
     * @return: the maximum average
     */
    double maxAverage(vector<int> &nums, int k) {
        // write your code here
        //sum: the sum of the last k elements ending with the (i-1)th element;
        double res = 0, sum = 0;
        int len = nums.size();
        stack<pair<double, double>> sk;
        for(int i=0; i<k; ++i) sum += (double) nums[i];
        res = sum/(double)k;
        sk.push({sum, k});
        for(int i=k; i<len; ++i) {
            sum += (double)nums[i] - (double)nums[i-k];
            vector<pair<double, double>> vs;
            while(sk.size()) {
                vs.push_back(sk.top());
                sk.pop();
            }
            for(auto &a : vs) {
                a.first += (double) nums[i];
                a.second += 1;
            }
            vs.push_back({sum, k});
            for(auto &a : vs) {
                auto t = a.first/a.second;
                if(res < t) res = t;
                if(sk.empty() || sk.top().first/sk.top().second > t) sk.push(a);
                else {
                    while(sk.size() && sk.top().first/sk.top().second <= t) sk.pop();
                    sk.push(a);
                }
            }
        }
        return res;
    }
};

Follow up:

This question can be modified a little bit: what is the maximum subarray sum with the size of the subarray no less than k?

The idea is very similar to the above one, but it is much easier to code, since we just need to save the largest sum from all previous subarray ending with the (i-1)th element.