Skip to content

Instantly share code, notes, and snippets.

@vetional
Created October 6, 2017 14:09
Show Gist options
  • Save vetional/63f86dbae4b7a6da52b0ec0d12d487bd to your computer and use it in GitHub Desktop.
Save vetional/63f86dbae4b7a6da52b0ec0d12d487bd to your computer and use it in GitHub Desktop.
Prefix sum (mashroom) created by vetional - https://repl.it/MNNH/0
'''
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment