NORTH, SOUTH, WEST, EAST = (0, -1), (0, 1), (-1, 0), (1, 0)
def max_water_heldover_mike(heights, display=None):
"""Algorithm from @Mike's answer.
https://ru.stackoverflow.com/a/727959/23044
"""
# fill initially upto the max height
max_height = max(max(row) for row in heights)
L = [[max_height] * len(row) for row in heights]
def lower_water_level_in_cell(x, y, level):
"""Lower the water level in the cell *x*, *y* down to *level* if possible."""
if not (0 <= y < len(heights) and 0 <= x < len(heights[y])):
return # the cell is outside of the island
if L[y][x] <= level: # water level in the cell is lower
return # the cell is already drained
# the water level can't be less than the cell height
level = max(level, heights[y][x])
L[y][x] = level # lower the water level
display and display(L)
for dx, dy in [NORTH, SOUTH, WEST, EAST]:
lower_water_level_in_cell(x + dx, y + dy, level)
# start with the island perimeter
for y in range(len(heights)):
lower_water_level_in_cell(x=0, y=y, level=0) # WEST side
lower_water_level_in_cell(x=len(heights[y]) - 1, y=y, level=0) # EAST
for x in range(len(heights[0])):
lower_water_level_in_cell(x=x, y=0, level=0) # NORTH side
for x in range(len(heights[len(heights) - 1])):
lower_water_level_in_cell(x=x, y=len(heights) - 1, level=0) # SOUTH
# find how much water is left
return sum((level - height)
for row, h_row in zip(L, heights)
for level, height in zip(row, h_row))
if __name__ == "__main__":
import sys
heights = [[int(n) for n in line.split()]
for line in sys.stdin if line.strip()]
print(max_water_heldover_mike(heights))
Tk9SVEgsIFNPVVRILCBXRVNULCBFQVNUID0gKDAsIC0xKSwgKDAsIDEpLCAoLTEsIDApLCAoMSwgMCkKCmRlZiBtYXhfd2F0ZXJfaGVsZG92ZXJfbWlrZShoZWlnaHRzLCBkaXNwbGF5PU5vbmUpOgogICAgIiIiQWxnb3JpdGhtIGZyb20gQE1pa2UncyBhbnN3ZXIuCgogICAgaHR0cHM6Ly9ydS5zdGFja292ZXJmbG93LmNvbS9hLzcyNzk1OS8yMzA0NAogICAgIiIiCiAgICAjIGZpbGwgaW5pdGlhbGx5IHVwdG8gdGhlIG1heCBoZWlnaHQKICAgIG1heF9oZWlnaHQgPSBtYXgobWF4KHJvdykgZm9yIHJvdyBpbiBoZWlnaHRzKQogICAgTCA9IFtbbWF4X2hlaWdodF0gKiBsZW4ocm93KSBmb3Igcm93IGluIGhlaWdodHNdCgogICAgZGVmIGxvd2VyX3dhdGVyX2xldmVsX2luX2NlbGwoeCwgeSwgbGV2ZWwpOgogICAgICAgICIiIkxvd2VyIHRoZSB3YXRlciBsZXZlbCBpbiB0aGUgY2VsbCAqeCosICp5KiBkb3duIHRvICpsZXZlbCogaWYgcG9zc2libGUuIiIiCiAgICAgICAgaWYgbm90ICgwIDw9IHkgPCBsZW4oaGVpZ2h0cykgYW5kIDAgPD0geCA8IGxlbihoZWlnaHRzW3ldKSk6CiAgICAgICAgICAgIHJldHVybiAgIyB0aGUgY2VsbCBpcyBvdXRzaWRlIG9mIHRoZSBpc2xhbmQKICAgICAgICBpZiBMW3ldW3hdIDw9IGxldmVsOiAgIyB3YXRlciBsZXZlbCBpbiB0aGUgY2VsbCBpcyBsb3dlcgogICAgICAgICAgICByZXR1cm4gICMgdGhlIGNlbGwgaXMgYWxyZWFkeSBkcmFpbmVkCiAgICAgICAgIyB0aGUgd2F0ZXIgbGV2ZWwgY2FuJ3QgYmUgbGVzcyB0aGFuIHRoZSBjZWxsIGhlaWdodAogICAgICAgIGxldmVsID0gbWF4KGxldmVsLCBoZWlnaHRzW3ldW3hdKQogICAgICAgIExbeV1beF0gPSBsZXZlbCAgIyBsb3dlciB0aGUgd2F0ZXIgbGV2ZWwKICAgICAgICBkaXNwbGF5IGFuZCBkaXNwbGF5KEwpCiAgICAgICAgZm9yIGR4LCBkeSBpbiBbTk9SVEgsIFNPVVRILCBXRVNULCBFQVNUXToKICAgICAgICAgICAgbG93ZXJfd2F0ZXJfbGV2ZWxfaW5fY2VsbCh4ICsgZHgsIHkgKyBkeSwgbGV2ZWwpCgogICAgIyBzdGFydCB3aXRoIHRoZSBpc2xhbmQgcGVyaW1ldGVyCiAgICBmb3IgeSBpbiByYW5nZShsZW4oaGVpZ2h0cykpOgogICAgICAgIGxvd2VyX3dhdGVyX2xldmVsX2luX2NlbGwoeD0wLCB5PXksIGxldmVsPTApICAjIFdFU1Qgc2lkZQogICAgICAgIGxvd2VyX3dhdGVyX2xldmVsX2luX2NlbGwoeD1sZW4oaGVpZ2h0c1t5XSkgLSAxLCB5PXksIGxldmVsPTApICAjIEVBU1QKICAgIGZvciB4IGluIHJhbmdlKGxlbihoZWlnaHRzWzBdKSk6CiAgICAgICAgbG93ZXJfd2F0ZXJfbGV2ZWxfaW5fY2VsbCh4PXgsIHk9MCwgbGV2ZWw9MCkgICMgTk9SVEggc2lkZQogICAgZm9yIHggaW4gcmFuZ2UobGVuKGhlaWdodHNbbGVuKGhlaWdodHMpIC0gMV0pKToKICAgICAgICBsb3dlcl93YXRlcl9sZXZlbF9pbl9jZWxsKHg9eCwgeT1sZW4oaGVpZ2h0cykgLSAxLCBsZXZlbD0wKSAgIyBTT1VUSAoKICAgICMgZmluZCBob3cgbXVjaCB3YXRlciBpcyBsZWZ0CiAgICByZXR1cm4gc3VtKChsZXZlbCAtIGhlaWdodCkKICAgICAgICAgICAgICAgZm9yIHJvdywgaF9yb3cgaW4gemlwKEwsIGhlaWdodHMpCiAgICAgICAgICAgICAgIGZvciBsZXZlbCwgaGVpZ2h0IGluIHppcChyb3csIGhfcm93KSkKCmlmIF9fbmFtZV9fID09ICJfX21haW5fXyI6CiAgICBpbXBvcnQgc3lzCiAgICBoZWlnaHRzID0gW1tpbnQobikgZm9yIG4gaW4gbGluZS5zcGxpdCgpXQogICAgICAgICAgICAgICAgICAgICBmb3IgbGluZSBpbiBzeXMuc3RkaW4gaWYgbGluZS5zdHJpcCgpXQogICAgcHJpbnQobWF4X3dhdGVyX2hlbGRvdmVyX21pa2UoaGVpZ2h0cykpCg==