python - prefix sum algorithm

You are not alone in considering the loop construction to be counter-intuitive, as I had to spend a few minutes on it as well. Here's what I figured out.

Now, the solution in the link you provided further details the optimal strategy is walking on path in such a way that one changes directions only once. In that manner, one is able to cover a range with left and right endpoints, which left_pos and right_pos seems to represent.

As to the particulars of the loops, instead of thinking of the loop in terms of the loop variables(i.e. p) it is easier to figure out what changes through the course of the loop, and how p is used. Otherwise, figuring out what is in those min and max expressions seems a bit too peculiar in the beginning.

For instance, in the first loop, instead of figuring out what that range represents, try how left_pos is affected by different values p gets. After a bit of thinking, one notices that left_pos changes in a manner complying to the possible left endpoints.

Specifically, when p == 0, left endpoint is the starting index(i.e. k) and when p is min(m, k), then it is either 0(i.e. if k < m) or (k - m). In the former case, that is as far as the left endpoint can go, as it would get out of the valid range of spots on the road. In the latter case, the number of moves prohibit any solution with a left_pos smaller than (k - m) since it is impossible to go from k to those indices in m moves.

The assignment made to right_pos in the first loop can be explained similarly. min statement includes (n-1), which is the rightmost legal index that can be reached and it serves to keep the right endpoint in the allowed limits. The inner max statement features k, as it is the least possible value for right_pos. (i.e. due to k being the starting point) It also has an expression (k + m - 2 * p). This expression represents the following process:

  • Go to left for p moves.
  • Change direction, and go to right for p moves to reach the starting point.
  • Go to right with the remaining (m - 2p) moves.

The second loop is just the reflection of this first loop, and you may explain it simply by adapting my explanation of the first loop.

As to your second question, I do not think it is common practice to shift the indices for prefix sum arrays. I typically use this method in online programming contests and my implementation of the prefix sum array you use in Python would be as follows.

def prefix_sums(A):
    n = len(A)
    P = [0] * n
    P[0] = A[0]
    for k in xrange(1, n):
        P[k] = P[k - 1] + A[k]
    return P

def count_total(P, x, y):
    return (P[y] - P[x - 1] if x > 0 else P[y])

The fundamental idea behind the implementation above is that, at P[x], we have the sum A[0] + A[1] + ... + A[x].


After reading the topic it was still hard to understand the idea, until i implemented naive solution(which is first in the codility document)

Hard to understand solution #2 simply imitates moving left and right and all these weird looking calculations only for getting left and right limits of the area(as you would really move inside it). So each iteration means one full cycle of using 6 steps.

If you move to the left and then to the right (p=0...M), you have

  • 0 steps left, 6 steps right(really 0 and 2 steps cause out of array border), so left border of area is at index 4 and right border is at index 6
  • 1 steps left, 5 steps right(really 1 and 3), so left border is at index 3 and right border is at index 6
  • 2 steps left, 4 steps right(really 2 and 4)...continue calculations

Here is my PHP version with oversimplified code and additional variables for easier understanding

function prefix_sums(array $a)
{
    $n = count($a);
    $p = array_fill(0, $n + 1, 0);
    for ($i = 1; $i <= $n; $i++) {
        $p[$i] = $p[$i - 1] + $a[$i - 1];
    }
    return $p;
}

function count_total($p, $x, $y)
{
    return $p[$y + 1] - $p[$x];
}

function mushrooms(array $a, int $k, int $m)
{
    $n = count($a) - 1;
    $max = 0;
    $sums = prefix_sums($a);
    //start  moving to the left and then the right
    for ($p = 0; $p < $m; $p++) {
        $stepsLeft = $p;
        $realStepsLeft = min($k, $stepsLeft);
        $leftBorder = $k - $realStepsLeft;

        $stepsRight = $m - $stepsLeft;
        $realStepsRight = min($n - $leftBorder, $stepsRight);
        $rightBorder = $leftBorder + $realStepsRight;

        $max = max($max, count_total($sums, $leftBorder, $rightBorder));
    }
    //moving to the right and then the left
    for ($p = 0; $p < $m; $p++) {
        $stepsRight = $p;
        $realStepsRight = min($p, $n - $k);
        $rightBorder = $k + $realStepsRight;

        $stepsLeft = $m - $stepsRight;
        $realStepsLeft = min(($k + $realStepsRight), $stepsLeft);
        $leftBorder = $rightBorder - $realStepsLeft;

        $max = max($max, count_total($sums, $leftBorder, $rightBorder));
    }
    return $max;
}

assert(ASSERT_EXCEPTION, 1);
assert(mushrooms([2, 3, 7, 5, 1, 3, 9], 4, 6) == 25);

echo 'Success';