Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- from collections import namedtuple
- S = int(input())
- elevation = [list(map(int, input().split())) for _ in range(S)]
- Point = namedtuple('Point', 'x y')
- valid_xy = [Point(x, y) for x in range(S) for y in range(S)]
- water = [[1 for x in range(S)] for y in range(S)]
- def neighbours(x, y):
- possible = (Point(x, y-1), Point(x-1, y),
- Point(x, y+1), Point(x+1, y))
- return (n for n in possible if n in valid_xy)
- def lowest_neighbour(x, y):
- lowest_point = None
- lowest_elev = elevation[x][y]
- for n in neighbours(x, y):
- if elevation[n.x][n.y] < lowest_elev:
- lowest_elev = elevation[n.x][n.y]
- lowest_point = Point(n.x, n.y)
- return lowest_point
- def drain(x, y):
- lowest = lowest_neighbour(x, y)
- if lowest:
- water[lowest.x][lowest.y] += water[x][y]
- water[x][y] = 0
- return True
- # Loop through all valid plots and drain them (This isn't elegant!)
- updated = True
- while updated:
- updated = False
- for p in valid_xy:
- if water[p.x][p.y]:
- if drain(*p):
- updated = True
- basins = [water[p.x][p.y] for p in valid_xy if water[p.x][p.y]]
- basins.sort(reverse=True)
- for x in basins: print(x, end=' ')
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement