Codility: Solution of Max Mushroom Picking Problem in Python

You are given a non-empty, zero-indexed array A of n (1 ¬ n ¬ 100 000) integers a0, a1, . . . , an−1 (0 ¬ ai ¬ 1 000). This array represents number of mushrooms growing on the consecutive spots along a road. You are also given integers k and m (0 ¬ k, m < n). A mushroom picker is at spot number k on the road and should perform m moves. In one move she moves to an adjacent spot. She collects all the mushrooms growing on spots she visits. The goal is to calculate the maximum number of mushrooms that the mushroom picker can collect in m moves.

Solution

def prefixSum(A):
    n = len(A)
    prefix = [0]*(n + 1)
    
    for k in range(1, n+1):
        prefix[k] = prefix[k - 1] + A[k - 1]
    print(prefix)
    return prefix

def maxMushrooms(A, k, m):
    prefix = prefixSum(A)
    l = len(A)
    org_k = k
    answer = 0
    
    start = k
    
    # case 1. Move forward
    end = min(k + m, l-1)
    pdiff = prefix[end+1] - prefix[start]
    print(pdiff, prefix[end+1], prefix[start])
    answer = max(answer, pdiff)
    
    # case 2 Move left
    end = k
    
    start = k - min(k, m)
    pdiff = prefix[end+1] - prefix[start]
    print(pdiff, prefix[end+1], prefix[start])
    answer = max(answer, pdiff)
    
    # case 3: Mix left & right
    
    # At position k you can go back least of both.
    # k = 4 m = 2, you can't move more than given moves i.e. 2
    # k= 1 m = 10 you can't move more than 1 move.
    max_ops = min(m, k) 
    
    while max_ops > 0:
        start = k - max_ops
        #print("start={}".format(start))
        
        print("remaining steps m={} consumes={}".format(m, (2 * max_ops)))
        # m - (2 * max_ops) can have negative value.
        # m = 6, max_ops = 4 diff = 6 - (2*4) this means
        # no operations remain to move right from position k.
        # So the end is k itself.
        end = min(k +  max(0, m - (2 * max_ops)), l-1)
        print("end={}".format(end))

        pdiff = prefix[end+1] - prefix[start]
        print(start, end+1, prefix[start], prefix[end+1])
        answer = max(answer, pdiff)
        max_ops -= 1

    print("answer={}".format(answer))
    
A = [2,3,1,7,5,1,3,9]
k = 4
m = 6 
#maxMushrooms(A, k, m)

A = [9,1,1,1,1,1,9]
k = 3
m = 4
#maxMushrooms(A, k, m)

A = [9,2,1,1,1,1,9]
k = 3
m = 10
print(A)
maxMushrooms(A, k, m)

Reference

Written with StackEdit.

Leave a Reply

Please log in using one of these methods to post your comment:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s

This site uses Akismet to reduce spam. Learn how your comment data is processed.

%d bloggers like this: