Prefix sum (mashroom) created by vetional - https://repl.it/MNNH/0

Prefix sum (mashroom).py

'''
Problem: 
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.
'''

import math

def prefix_sum(A):
  res = [0]
  sum = 0
  for x in A:
    sum += x
    res.append(sum)
  return res

# use prefix_sum to compute slice sum
def slice_sum(ps, p, q):
  q = min(len(ps) - 2, q)
  p = max(0, p)
  return ps[q+1] - ps[p]
  
def solution(A, s, m):
  ps = prefix_sum(A)
  max_collection = 0
  if m < 2 * max(len(A) - s, s):
    # extreme cases
    ranges = [(s, s + m), (s - m, s), (s + 1, s + m + 1), (s - m - 1, s -1)]
    for r in ranges:
      collection = slice_sum(ps, r[0], r[1])
      if collection > max_collection:
        max_collection = collection
    # window partially on each side
    for x in reversed(xrange(m / 2)):
      collection = slice_sum(ps, s - x, s + (m - 2 * (s - x)))
      if collection > max_collection:
        max_collection = collection
  else:
    slice_sum(ps, 0, len(A))
  return max_collection
    
    

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

print solution(A, s, m)