matrix-land-description.txt

Let's say we have a 4*16 matrix looks like this:

[ 4  5  3  0 -3  0  3  1 -4  3  8 -4 -1  9  1  8]
[ 2 -1 -6  8  8 -7 -5  1  0 -4 -5 -5  2 -7 -8  6]
[-7 -8 -7 -6 -5 -1 -3  3  6 -7  4 -7 -5  0 -5  1]
[-9  6 -4 -1  7 -7 -7  5  8  8  3 -6  3 -2  0 -1]
We want to calculate the highest value path to each square like this:

[33 33 33 33 33 33 33 33 33 33 33 33 33 33 33 33]
[44 44 44 49 49 42 38 38 38 34 29 30 35 28 31 39]
[37 36 37 43 49 49 49 49 49 46 46 39 34 35 35 40]
[58 67 67 67 67 67 67 73 73 73 73 70 70 68 68 67]
The answer of course will be the highest value on the bottom row, 73 in this case.

This is done line by line. For each square, there are three parameters that define the highest value path:

How far left the path goes on that line.
How far right the path goes on that line.
The square on the line above you drop down from.
We can brute force this, but it is O(n*m^4).

We can reduce that to O(n*m^3) if we note that how far the path goes to the left is largely independent of how far it goes to the right, except for which side we drop down from the line above. We can therefore find the best way to drop down from either side, then how far we should go in either direction if we don't drop down that side.

We can further reduce it to O(n*m^2) because finding the best square to drop down from is a simple maximum that we can calculate while figuring out how far we should go. But we really need O(n*m).

Kadane's algorithm is the key to this. Kadane's algorithm finds the maximum subarray in an array. It looks like this:

max_subarray = arr[0]
subtotal = 0
for value in arr:
    subtotal = max(value, subtotal + value)
    max_subarray = max(subtotal, max_subarray)
The subtotal gives the maximum subarray that stops at a particular square. If we proceed rightwards through a line on the grid, we can determine how far left the path should go from each square. If we proceed leftwards we can also determine how far right to go from each square. And this is O(m).

Accounting for the fact that we must drop down from a square above can be done with the same principle.

import sys
from array import array

def matrixLand(grid):
    width = len(grid[0])
    prevline = array('l',[0]) * width

    for line in grid:
        going_left = array('l',[0]) * width
        going_right = array('l',[0]) * width

        top, notop = -2000000000, 0
        for i in range(width):
            notop = max(0,notop)
            going_right[i] += notop
            notop += line[i]
            top = max(top+line[i], prevline[i]+notop)
            going_left[i] += top

        top, notop = -2000000000, 0
        for i in range(width-1,-1,-1):
            notop = max(0,notop)
            going_left[i] += notop
            notop += line[i]
            top = max(top+line[i], prevline[i]+notop)
            going_right[i] += top

        for i in range(width):
            prevline[i] = max(going_left[i], going_right[i])

    return max(prevline)


n, m = [int(x) for x in input().split()]

if m == 1:
    print(sum(int(x) for x in sys.stdin))
else:
    grid = []
    for _ in range(n):
        grid.append(array('l',(int(x) for x in sys.stdin.readline().split(' '))))
    result = matrixLand(grid)
    print(result)
You may notice that I have written special case code for the m = 1 case. TC#19 is n = 4 million, m = 1, and the per line overhead is just too much for PyPy. Even the special case code, using sys.stdin instead of input(), takes 3 seconds. TC#20 is n = 1, m = 4 million and forces us to use arrays instead of lists to avoid running out of memory.