Python, Straightforward with Explanation

  • 5

    We binary search on the answer. Let P[i] = A[0] + A[1] + ... + A[i-1], the i-th prefix sum under A.

    Let's focus our attention on possible(x), a function that is true iff it is possible to have an average of at least x. Consider the elements B = [a-x for a in A] with corresponding prefix sum Q[i] = P[i] - i*x under B.

    We want to know if there is some >= K length subarray in B with average at least zero. Suppose the subarray is B[i] + B[i+1] + ... + B[j] = Q[j+1] - Q[i]. To check whether this quantity is positive, for any j, and any i <= j - K + 1, we should check whether Q[j+1] >= min_{i <= j-K+1} Q[i]. Keeping a running minimum m of this array Q, we can check this in linear time.

    Unfortunately, the time constraint on Python solutions is fairly tight, so we need another trick to avoid TLE. If a segment has the biggest average and we break it into two pieces, one of its pieces also has at least the same average. When the length is >= 2*K, we can split it into pieces of at least length K, with the largest such piece being less than length 2*K.

    Thus, we only need to check segments of length K <= L < 2*K to find an instance of the maximum average. When K is small, this admits an O(NK) solution that we use instead. Our solution in that case is identical to Maximum Average Subarray I, repeated K times.

    def findMaxAverage(self, A, K):
        N = len(A)
        P = [0]
        for x in A:
            P.append(P[-1] + x)
        if K < 100:
            ans = float('-inf')
            for k in xrange(K, min(2*K, N+1)):
                best_sum = max(P[i+k] - P[i] for i in xrange(N-k+1))
                ans = max(ans, best_sum / float(k))
            return ans
        def possible(x):
            m = P[0]
            for i, v in enumerate(P):
                m = min(m, v-i*x)
                if i+K == len(P): break
                if P[i+K] - (i+K)*x >= m:
                    return True
            return False
        lo, hi = min(A), max(A)
        while hi - lo > .00001:
            mi = (lo + hi) / 2.0
            if possible(mi):
                lo = mi
                hi = mi
        return lo

  • 0

    @awice Great idea! I got both but didn't piece them together.

    A small question: how do you get the threshold for k 100? Is there any reason or just keep trying?

  • 0

    @awice said in Python, Straightforward with Explanation:

    with the smallest such piece being less than length 2*K.

    Thanks for the great idea!
    One question is that "with the smallest such piece being less than length 2K", do you actually mean "with the longest such piece being less than length 2K."? Because later on you wrote "we only need to check segments of length K <= L < 2*K to find an instance of the maximum average."
    Just want to confirm my understanding.

  • 0

    @dzliu2007-gmail-com smallest, yes. I edited my post now.

    @flamesofmoon It was the first condition I tried and it worked when the other failed. We can try to estimate the number of steps in the other method as log_2(10^9) (N-K)K.

  • 0

    Is the running time O(n * log2(10^9)) = O(n)?

  • 0

    @awice I have a question that does not have to do with the problem, regarding the line lo, hi = min(A), max(A). I know that the min and max functions are written in native C and run much faster than using a loop in Python. But, do they run so much faster that iterating twice through the array (as lo, hi = min(A), max(A) does) is faster than iterating once with two variables to keep the answer?

  • 1

    @david120 Is this what you meant?

    from random import randint
    from time import time
    A = [randint(0, 123456) for _ in xrange(1000000)]
    def solve1(A):
        return min(A), max(A)
    def solve2(A):
        lo = float('inf')
        hi = float('-inf')
        for x in A:
            lo = min(lo, x)
            hi = max(hi, x)
        return lo, hi
    for fn in (solve1, solve2, solve1, solve2):
        tt = time()
        for i in xrange(10):
            ans = fn(A)
        print time() - tt

  • 0

    @awice Ah yes I should have just done the test myself instead of asking. It is pretty amazing that running min and max is almost an entire order of magnitude faster than finding the answers using a Python loop.

Log in to reply

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.