fork(5) download
  1. from collections import namedtuple
  2. from queue import PriorityQueue
  3.  
  4. NORTH, SOUTH, WEST, EAST = (0, -1), (0, 1), (-1, 0), (1, 0)
  5.  
  6. Cell = namedtuple('Cell', 'level x y')
  7.  
  8. def max_water_heldover_minheap(heights, display=None):
  9. """from http://w...content-available-to-author-only...u.com/p/540ee7ec725b"""
  10. q = PriorityQueue()
  11. nrows = len(heights)
  12. ncolumns = len(heights[0])
  13. seen = [[False] * ncolumns for _ in range(nrows)]
  14.  
  15. # enqueue cells on the perimeter of the island
  16. for y in range(nrows):
  17. seen[y][0] = True
  18. q.put(Cell(heights[y][0], 0, y)) # WEST side
  19. seen[y][ncolumns - 1] = True
  20. q.put(Cell(heights[y][ncolumns - 1], ncolumns - 1, y)) # EAST
  21.  
  22. for x in range(ncolumns):
  23. seen[0][x] = True
  24. q.put(Cell(heights[0][x], x, 0)) # NORTH side
  25. seen[nrows - 1][x] = True
  26. q.put(Cell(heights[nrows - 1][x], x, nrows - 1)) # SOUTH
  27.  
  28. # visit all cells once starting with cells with a minimum water level
  29. total = 0
  30. while not q.empty():
  31. cell = q.get()
  32. for dx, dy in [NORTH, SOUTH, WEST, EAST]:
  33. x = cell.x + dx
  34. y = cell.y + dy
  35. if 0 <= y < nrows and 0 <= x < ncolumns and not seen[y][x]:
  36. seen[y][x] = True
  37. q.put(Cell(max(cell.level, heights[y][x]), x, y))
  38. total += max(0, cell.level - heights[y][x])
  39.  
  40. return total
  41.  
  42. if __name__ == "__main__":
  43. import sys
  44. heights = [[int(n) for n in line.split()]
  45. for line in sys.stdin if line.strip()]
  46. print(max_water_heldover_minheap(heights))
  47.  
Success #stdin #stdout 0.03s 28632KB
stdin
5 3 4 5
6 2 1 4
3 1 1 4
8 5 5 3
8 3 1 6
8 4 5 8
stdout
11