Mushroom Picker (A problem on Prefix Sums)

Problem

You are given a non-empty, zero-indexed array A of n (1 <= n <= 100000) integers a0, a1,..., an−1 (0 <= ai <= 1000). 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.

For example, consider array A such that:

A[0]=2, A[1]=3, A[2]=7, A[3]=5, A[4]=1, A[5]=3, A[6]=9

The mushroom picker starts at spot k = 4 and should perform m = 6 moves. She might

move to spots 3, 2, 3, 4, 5, 6 and thereby collect 1 + 5 + 7 + 3 + 9 = 25 mushrooms. This is the

maximal number of mushrooms she can collect.

Code

# Counting prefix sums
  def prefix_sums(A):
      n = len(A)
      P = [0] * (n + 1)
      for k in xrange(1, n + 1):
          P[k] = P[k - 1] + A[k - 1]
      return P
  
  # Total of one slice
  def count_total(P, x, y):
      return P[y + 1] - P[x]
  # Mushroom picker — O(n + m)
  def mushrooms(A, k, m):
      n = len(A)
      result = 0
      pref = prefix_sums(A)
  
  # When we first take p steps on the left and return from their back in right direction.
      for p in xrange(min(m, k) + 1):
          left_pos = k - p
          right_pos = min(n - 1, max(k, k + m - 2*p))
          result = max(result, count_total(pref, left_pos, right_pos)) 
  
  # When we first take p steps on the right and return from their back in left direction.
      for p in xrange(min(m + 1, n - k)):
          right_pos = k + p
          left_pos = max(0, min(k, k - (m - 2*p)))
          result=max(result, count_total(pref, left_pos, right_pos)) 
      return result

Explanation (Part 1)

# When we first take p steps on the left and return from their back in right direction.
    for p in xrange(min(m, k) + 1):
        left_pos = k - p
        right_pos = min(n - 1, max(k, k + m - 2*p))
        result = max(result, count_total(pref, left_pos, right_pos)) 

Expression “for p in xrange(min(m, k) + 1)” has min(m, k) because we can only go m steps to the left or k steps to the left whichever number is minimum.

After going k steps, we cannot go past 0th position.

Or after taking m steps, we cannot take another step.

Expression “right_pos = min(n - 1, max(k, k + m – 2*p))” has ‘k+m-2*p’ because after taking p steps to the left, we are returning p steps back to position k, hence 2p.

And then number of steps left is: m – 2p

Then right position is identified by: m-2p steps after k ie, k+m-2p

Expression “right_pos = min(n - 1, max(k, k + m – 2*p))” has max(k, k + m – 2*p) because what if we take all the m steps to left and value of p becomes m.

Then there are no steps left to take to the right and right position is simply identified by k.

Similarly, for any value of p > m/2 (as in 0.6m or 0.75m), we would get same right position as k.

Explanation (Part 2)

# When we first take p steps on the right and return from their back in left direction.
    for p in xrange(min(m + 1, n - k)):
        right_pos = k + p
        left_pos = max(0, min(k, k - (m - 2*p)))
        result = max(result, count_total(pref, left_pos, right_pos))

Here in the expression “for p in xrange(min(m + 1, n – k))”, we have m+1 and not just m because xrange goes till m when passed the argument m+1. Similarly, we have n-k and not (n-1)-k because xrange goes till (n-1)-k when passed the argument n-k.

And we take minimum of m or (n-1)-k because that’s is the possible number of steps we can take on the right.

Here in the expression “left_pos = max(0, min(k, k - (m – 2*p)))”, we have k-(m-2p) because we take p steps on the right, then take those p steps back to k (on the left now).

Number of steps yet to take: m-2p

Left position would be identified by: k-(m-2p)

If p takes value m, which means we have taken m steps on the right, then we have no steps remaining to take to the left and left position is identified by k itself. (This logic is valid for any value of p > m/2)

Side Note

The way we calculating prefix sums and sums of subsections is without using the Python built-ins. But if we were to use Python built-ins, the code would look somewhat like this:

def get_mushrooms_sum(A, start, end):
  return sum(A[start : end + 1])