fork download
copy 
  1. # -*- coding: utf-8 -*-
  2. import copy
  3. from collections import defaultdict
  4. from itertools import product
  5. import heapq
  6.  
  7. FALL = '|'
  8. RIVER = '~'
  9. POOL = 'N'
  10.  
  11. SOURCE = 'x'
  12. SOURCE_AND_FALL = '*'
  13. SOURCE_AND_RIVER = 'X'
  14. SOURCE_AND_POOL = '%'
  15.  
  16. BLOCK = '#'
  17. LRAMP = '/'
  18. RRAMP = '\\'
  19. SPACE = ' '
  20.  
  21. LRAMP_AND_FALL = 'd'
  22. RRAMP_AND_FALL = 'b'
  23. LRAMP_AND_RIVER = '}'
  24. RRAMP_AND_RIVER = '{'
  25. LRAMP_AND_POOL = ']'
  26. RRAMP_AND_POOL = '['
  27.  
  28. DOWN = (1,0)
  29. LEFT = (0,-1)
  30. RIGHT = (0,1)
  31. UP = (-1,0)
  32. ORDINALS = [DOWN,LEFT,RIGHT,UP]
  33.  
  34. gridStr = """\
  35.    x      
  36. #        #
  37. #        #
  38. #        #
  39. #### #####
  40. #### #####
  41. #### #####
  42. ##########"""
  43. gridStr = """\
  44.   x   
  45. #   / 
  46. ##### 
  47. #   # 
  48. # #  /
  49. ###\/#"""
  50. gridStr = """\
  51.   x  
  52. #  / 
  53. #####"""
  54. gridStr = """\
  55.                     x  
  56.  
  57.  
  58. #                   # #
  59. #                   # #
  60. #                   # #
  61. #                   # #
  62. #                   # #
  63. #######################"""
  64. gridStr = """\
  65.  x #  
  66. #  # #
  67. #  # #
  68. #  # #
  69. #    #
  70. ######"""
  71.  
  72. # The capacity of a ramp cell.
  73. HALF_CAPACITY = 50
  74. # The capacity of an empty cell.
  75. MAX_CAPACITY = 2*HALF_CAPACITY
  76. # The rate of flow in from a source.
  77. SOURCE_RATE = HALF_CAPACITY
  78. # The maximum rate of flow between two cells.
  79. MAX_FLOW = MAX_CAPACITY
  80.  
  81. SPACE_THRESHHOLD = HALF_CAPACITY / 2
  82. RIVER_THRESHHOLD = HALF_CAPACITY + HALF_CAPACITY / 2
  83.  
  84. def isSource(cell):
  85.     return cell in [SOURCE, SOURCE_AND_FALL, SOURCE_AND_RIVER, SOURCE_AND_POOL]
  86.  
  87. def isLRamp(cell):
  88.     return cell in [LRAMP, LRAMP_AND_FALL, LRAMP_AND_RIVER, LRAMP_AND_POOL]
  89. def isRRamp(cell):
  90.     return cell in [RRAMP, RRAMP_AND_FALL, RRAMP_AND_RIVER, RRAMP_AND_POOL]
  91. def isRamp(cell):
  92.     return isLRamp(cell) or isRRamp(cell)
  93.  
  94. def findBackground(cell):
  95.     if cell in [SPACE, FALL, RIVER, POOL, SOURCE]:
  96.         return SPACE
  97.     elif cell in [BLOCK]:
  98.         return BLOCK
  99.     elif isLRamp(cell):
  100.         return LRAMP
  101.     elif isRRamp(cell):
  102.         return RRAMP
  103.     raise Exception
  104.  
  105. class WaterType:
  106.     FALL = 1
  107.     SETTLED = 2
  108.  
  109. class Cell:
  110.     def __init__(self, background, isSource):
  111.         self.background = background
  112.         self.isSource = isSource
  113.         self.water = 0
  114.         self.waterType = WaterType.FALL
  115.     def getCapacity(self):
  116.         return (0 if self.background == BLOCK else (HALF_CAPACITY if isRamp(self.background) else MAX_CAPACITY))
  117.     def copyCell(self, cell):
  118.         self.background = cell.background
  119.         self.isSource = cell.isSource
  120.         # Include any water already here
  121.         self.water += cell.water
  122.         self.waterType = cell.waterType
  123.     def updateType(self, cellBelow):
  124.         if self.water == 0:
  125.             self.waterType = WaterType.SETTLED
  126.         elif self.water == self.getCapacity():
  127.             # If a completely saturated block of water is falling it should still look full.
  128.             self.waterType = WaterType.SETTLED
  129.         elif cellBelow.water < cellBelow.getCapacity():
  130.             self.waterType = WaterType.FALL
  131.         else:
  132.             self.waterType = WaterType.SETTLED
  133.     def __str__(self):
  134.         if self.background == SPACE and not self.isSource and self.water <= SPACE_THRESHHOLD:
  135.             return SPACE
  136.         elif self.background == SPACE and not self.isSource and self.waterType == WaterType.FALL:
  137.             return FALL
  138.         elif self.background == SPACE and not self.isSource and self.water <= RIVER_THRESHHOLD:
  139.             return RIVER
  140.         elif self.background == SPACE and not self.isSource:
  141.             return POOL
  142.         elif self.background == SPACE and self.water <= SPACE_THRESHHOLD:
  143.             return SOURCE
  144.         elif self.background == SPACE and self.waterType == WaterType.FALL:
  145.             return SOURCE_AND_FALL
  146.         elif self.background == SPACE and self.water <= RIVER_THRESHHOLD:
  147.             return SOURCE_AND_RIVER
  148.         elif self.background == SPACE:
  149.             return SOURCE_AND_POOL
  150.         elif self.background == BLOCK:
  151.             return BLOCK
  152.         elif isRamp(self.background) and self.isSource:
  153.             raise Exception
  154.         elif self.background == LRAMP and self.water <= SPACE_THRESHHOLD / 2:
  155.             return LRAMP
  156.         elif self.background == LRAMP and self.waterType == WaterType.FALL:
  157.             return LRAMP_AND_FALL
  158.         elif self.background == LRAMP and self.water <= RIVER_THRESHHOLD / 2:
  159.             return LRAMP_AND_RIVER
  160.         elif self.background == LRAMP:
  161.             return LRAMP_AND_POOL
  162.         elif self.water <= SPACE_THRESHHOLD / 2:
  163.             return RRAMP
  164.         elif self.waterType == WaterType.FALL:
  165.             return RRAMP_AND_FALL
  166.         elif self.water <= RIVER_THRESHHOLD / 2:
  167.             return RRAMP_AND_RIVER
  168.         else:
  169.             return RRAMP_AND_POOL
  170.  
  171. gridCells = map(lambda x: list(x), gridStr.split('\n'))
  172. height = len(gridCells)
  173. width = len(gridCells[0])
  174. grid = defaultdict(lambda : Cell(SPACE, False))
  175. for i in range(height):
  176.     for j in range(width):
  177.         grid[i,j] = Cell(findBackground(gridCells[i][j]), isSource(gridCells[i][j]))
  178.  
  179. def p(coords, offset):
  180.     """Add two tuples together.
  181.  
  182.     >>> p((2, 3), (-1, 0))
  183.     (1, 3)
  184.     """
  185.     return tuple(map(sum, zip(coords, offset)))
  186.  
  187. def adjacent(i,j,includeUp):
  188.     ret = []
  189.     for offset in (ORDINALS if includeUp else ORDINALS[:-1]):
  190.         ret.append(p((i,j), offset))
  191.     return ret
  192.  
  193. CONNECTION_TO_SHAPES = {DOWN: [SPACE, LRAMP, RRAMP], LEFT: [SPACE, RRAMP], RIGHT: [SPACE, LRAMP], UP: [SPACE]}
  194.  
  195. def adjoint(grid, i, j, includeUp, includeDown):
  196.     """Find all cells connected to this one (if this is a block then nothing connects to it).
  197.  
  198.     >>> grid = defaultdict(lambda : Cell(SPACE, False))
  199.     >>> adjoint(grid, 10, 1, True, True)
  200.     [(11, 1), (10, 0), (10, 2), (9, 1)]
  201.  
  202.     >>> adjoint(grid, 10, 1, False, True)
  203.     [(11, 1), (10, 0), (10, 2)]
  204.  
  205.     >>> adjoint(grid, 10, 1, False, False)
  206.     [(10, 0), (10, 2)]
  207.  
  208.     >>> grid[10, 1].background = LRAMP
  209.     >>> adjoint(grid, 10, 1, True, True)
  210.     [(10, 0), (9, 1)]
  211.  
  212.     >>> grid[10, 0].background = LRAMP
  213.     >>> grid[9, 1].background = RRAMP
  214.     >>> adjoint(grid, 10, 1, True, True)
  215.     []
  216.     """
  217.     ret = []
  218.     cell = grid[i, j]
  219.     if cell.background == BLOCK:
  220.         adj = []
  221.     elif cell.background == LRAMP:
  222.         adj = [LEFT, UP]
  223.     elif cell.background == RRAMP:
  224.         adj = [RIGHT, UP]
  225.     elif cell.background == SPACE:
  226.         adj = list(ORDINALS)
  227.     if not includeUp and UP in adj:
  228.         adj.remove(UP)
  229.     if not includeDown and DOWN in adj:
  230.         adj.remove(DOWN)
  231.     for offset in adj:
  232.         neighbourCoords = p((i,j), offset)
  233.         neighbour = grid[neighbourCoords]
  234.         if neighbour.background in CONNECTION_TO_SHAPES[offset]:
  235.             ret.append(neighbourCoords)
  236.     return ret
  237.  
  238. def distance(loc1, loc2):
  239.     """Calculate the manhatten distance between two locations.
  240.  
  241.     >>> distance((5,2), (3, 1))
  242.     3
  243.     """
  244.     row1, col1 = loc1
  245.     row2, col2 = loc2
  246.     d_col = abs(col1 - col2)
  247.     d_row = abs(row1 - row2)
  248.     return d_row + d_col
  249.  
  250. def aStarRegion(start, end, region, edge):
  251.     """The AStar algorithm to find the shortest path between two cells using the region as the nodes, and the edges in edge.
  252.  
  253.     >>> grid = defaultdict(lambda : Cell(SPACE, False))
  254.     >>> edge = defaultdict(dict)
  255.     >>> edge[1,0] = {(1,1): 1}
  256.     >>> edge[1,1] = {(2,1): 1}
  257.     >>> aStarRegion((1,0), (2,1), [(1,0),(1,1),(2,1)], edge)
  258.     [(1, 0), (1, 1), (2, 1)]
  259.     """
  260.     openSet = set()
  261.     openHeap = []
  262.     closedSet = set()
  263.     parent = {}
  264.  
  265.     def retracePath(c, edge):
  266.         path = [c]
  267.         while c in parent:
  268.             c = parent[c]
  269.             path.append(c)
  270.         path.reverse()
  271.         return path
  272.  
  273.     current = start
  274.     openSet.add(current)
  275.     openHeap.append((0,current))
  276.     while openSet:
  277.         current = heapq.heappop(openHeap)[1]
  278.         if current == end:
  279.             return retracePath(current, edge)
  280.         openSet.remove(current)
  281.         closedSet.add(current)
  282.         for tile, remainingFlow in edge[current].items():
  283.             if tile not in closedSet:
  284.                 H = distance(tile, end)*10
  285.                 if tile not in openSet:
  286.                     openSet.add(tile)
  287.                     heapq.heappush(openHeap, (H,tile))
  288.                 parent[tile] = current
  289.     return []
  290.  
  291. def findWaterRegions(grid, height, width):
  292.     """Find water regions. Return a dict keyed on the top left (top then left) cell of each region.
  293.  
  294.     >>> grid = defaultdict(lambda : Cell(SPACE, False))
  295.     >>> waterRegions, outlets, edges = findWaterRegions(grid, 4, 5)
  296.     >>> waterRegions
  297.     {}
  298.     >>> outlets
  299.     {}
  300.     >>> dict(edges)
  301.     {}
  302.  
  303.     >>> grid[2,2].water = MAX_CAPACITY
  304.     >>> grid[2,3].water = MAX_CAPACITY
  305.     >>> grid[3,3].water = MAX_CAPACITY
  306.     >>> waterRegions, outlets, edges = findWaterRegions(grid, 4, 5)
  307.     >>> waterRegions
  308.     {(2, 2): [(2, 2), (2, 3), (3, 3)]}
  309.     >>> outlets
  310.     {(2, 2): [(3, 2), (2, 1), (1, 2), (2, 4), (1, 3), (4, 3), (3, 2), (3, 4)]}
  311.     >>> dict(edges)
  312.     {(2, 2): defaultdict(<type 'dict'>, {(2, 3): {(2, 4): 100, (1, 3): 100, (3, 3): 100, (2, 2): 100}, (3, 3): {(2, 3): 100, (3, 2): 100, (3, 4): 100, (4, 3): 100}, (2, 2): {(1, 2): 100, (3, 2): 100, (2, 3): 100, (2, 1): 100}})}
  313.  
  314.     >>> grid[2,3].background = LRAMP
  315.     >>> grid[2,3].water = HALF_CAPACITY
  316.     >>> waterRegions, outlets, edges = findWaterRegions(grid, 4, 5)
  317.     >>> waterRegions
  318.     {(3, 3): [(3, 3)], (2, 2): [(2, 2), (2, 3)]}
  319.     >>> outlets
  320.     {(3, 3): [(4, 3), (3, 2), (3, 4)], (2, 2): [(3, 2), (2, 1), (1, 2), (1, 3)]}
  321.     >>> dict(edges)
  322.     {(3, 3): defaultdict(<type 'dict'>, {(3, 3): {(3, 2): 100, (3, 4): 100, (4, 3): 100}}), (2, 2): defaultdict(<type 'dict'>, {(2, 3): {(1, 3): 100, (2, 2): 100}, (2, 2): {(1, 2): 100, (3, 2): 100, (2, 3): 100, (2, 1): 100}})}
  323.  
  324.     >>> grid = defaultdict(lambda : Cell(SPACE, False))
  325.     >>> grid[2,3].background = LRAMP
  326.     >>> grid[2,3].water = HALF_CAPACITY - 1
  327.     >>> waterRegions, outlets, edges = findWaterRegions(grid, 4, 5)
  328.     >>> waterRegions
  329.     {(2, 3): [(2, 3)]}
  330.     >>> outlets
  331.     {(2, 3): [(2, 3), (2, 2)]}
  332.     >>> dict(edges)
  333.     {(2, 3): defaultdict(<type 'dict'>, {(2, 3): {(2, 2): 100}})}
  334.  
  335.     >>> grid = defaultdict(lambda : Cell(SPACE, False))
  336.     >>> grid[1,0].background = BLOCK
  337.     >>> grid[1,4].background = BLOCK
  338.     >>> grid[2,0].background = BLOCK
  339.     >>> grid[2,4].background = BLOCK
  340.     >>> grid[3,1].background = RRAMP
  341.     >>> grid[4,2].background = BLOCK
  342.     >>> grid[2,2].background = BLOCK
  343.     >>> grid[3,3].background = LRAMP
  344.     >>> grid[3,1].water = HALF_CAPACITY
  345.     >>> grid[3,2].water = MAX_CAPACITY
  346.     >>> grid[3,3].water = HALF_CAPACITY
  347.     >>> waterRegions, outlets, edges = findWaterRegions(grid, 4, 5)
  348.     >>> waterRegions
  349.     {(3, 1): [(3, 1), (3, 2), (3, 3)]}
  350.     >>> outlets
  351.     {(3, 1): [(2, 1), (2, 3)]}
  352.     >>> dict(edges)
  353.     {(3, 1): defaultdict(<type 'dict'>, {(3, 2): {(3, 1): 100, (3, 3): 100}, (3, 1): {(3, 2): 100, (2, 1): 100}, (3, 3): {(3, 2): 100, (2, 3): 100}})}
  354.  
  355.     >>> grid = defaultdict(lambda : Cell(SPACE, False))
  356.     >>> grid[0,0].water = MAX_CAPACITY - 1
  357.     >>> grid[1,0].water = MAX_CAPACITY - 1
  358.     >>> waterRegions, outlets, edges = findWaterRegions(grid, 4, 5)
  359.     >>> waterRegions
  360.     {(1, 0): [(1, 0)], (0, 0): [(0, 0)]}
  361.     >>> outlets
  362.     {(1, 0): [(1, 0), (2, 0), (1, -1), (1, 1)], (0, 0): [(0, 0), (1, 0), (0, -1), (0, 1)]}
  363.     >>> dict(edges)
  364.     {(1, 0): defaultdict(<type 'dict'>, {(1, 0): {(2, 0): 100, (1, -1): 100, (1, 1): 100}}), (0, 0): defaultdict(<type 'dict'>, {(0, 0): {(0, 1): 100, (0, -1): 100, (1, 0): 100}})}
  365.     """
  366.     waterRegions = {}
  367.     outlets = {}
  368.     edges = defaultdict(lambda : defaultdict(dict))
  369.     explored = []
  370.     for i, j in product(range(height), range(width)):
  371.         if (i,j) not in explored:
  372.             cell = grid[i,j]
  373.             if cell.water > 0:
  374.                 waterRegions[i,j] = [(i,j)]
  375.                 outlets[i,j] = ([] if grid[i,j].water == grid[i,j].getCapacity() else [(i,j)])
  376.                 toExplore = [(i,j)]
  377.                 while len(toExplore) > 0:
  378.                     hereX, hereY = toExplore.pop()
  379.                     explored.append((hereX, hereY))
  380.                     includeUp = (grid[hereX,hereY].water == grid[hereX,hereY].getCapacity())
  381.                     includeDown = (grid[hereX+1,hereY].water == grid[hereX+1,hereY].getCapacity())
  382.                     for x, y in adjoint(grid, hereX, hereY, includeUp, includeDown):
  383.                         if grid[x,y].water != 0 and (x,y) not in explored:
  384.                             toExplore.append((x,y))
  385.                             waterRegions[i,j].append((x,y))
  386.                     for x, y in adjoint(grid, hereX, hereY, includeUp, True):
  387.                         edges[i,j][hereX,hereY][x,y] = MAX_FLOW
  388.                         if grid[x,y].water != grid[x,y].getCapacity():
  389.                             outlets[i,j].append((x,y))
  390.     return waterRegions, outlets, edges
  391.  
  392. def findHeight(grid, point):
  393.     return (height - point[0]) * MAX_CAPACITY + grid[point].water
  394.  
  395. def findHighestPoint(grid, region):
  396.     highestPoint = region[0]
  397.     highestHeight = findHeight(grid, highestPoint)
  398.     for point in region:
  399.         height = findHeight(grid, point)
  400.         if height > highestHeight:
  401.             highestHeight = height
  402.             highestPoint = point
  403.     return highestPoint
  404.  
  405. def findLowestPoint(grid, region):
  406.     lowestPoint = region[0]
  407.     lowestHeight = findHeight(grid, lowestPoint)
  408.     for point in region:
  409.         height = findHeight(grid, point)
  410.         if height < lowestHeight:
  411.             lowestHeight = height
  412.             lowestPoint = point
  413.     return lowestPoint
  414.  
  415. def moveWater(grid, start, end, region, edge):
  416.     """Move a unit of water from start to end. Return True if successful."""
  417.     path = aStarRegion(start, end, region, edge)
  418.     if len(path) == 0:
  419.         return False
  420.  
  421.     grid[start].water -= 1
  422.     stepStart = start
  423.     for stepEnd in path[1:]:
  424.         edge[stepStart][stepEnd] -= 1
  425.         if edge[stepStart][stepEnd] == 0:
  426.             edge[stepStart].pop(stepEnd)
  427.             if len(edge[stepStart]) == 0:
  428.                 edge.pop(stepStart)
  429.         stepStart = stepEnd
  430.     grid[end].water += 1
  431.  
  432.     return True
  433.  
  434. def displayGrid(grid):
  435.     print
  436.     for i in range(height):
  437.         row = ''
  438.         for j in range(width):
  439.             grid[i,j].updateType(grid[i+1,j])
  440.             row += str(grid[i,j])
  441.         row += ' '
  442.         for j in range(width):
  443.             row += ('%x'%(grid[i,j].water / 10) if grid[i,j].water < 160 else '+')
  444.         print row
  445.  
  446. def main(grid):
  447.     t = 0
  448.     while t < 20:
  449.         displayGrid(grid)
  450.         newGrid = defaultdict(lambda : Cell(SPACE, False))
  451.         for i, j in product(range(height), range(width)):
  452.             cell = grid[i,j]
  453.             newCell = newGrid[i,j]
  454.             newCell.copyCell(cell)
  455.             if cell.isSource:
  456.                 newCell.water += SOURCE_RATE
  457.         grid = newGrid
  458.  
  459.         waterRegions, outlets, edges = findWaterRegions(grid, height, width)
  460.  
  461.         movedSomething = True
  462.         while movedSomething:
  463.             movedSomething = False
  464.             for regionKey in waterRegions.keys():
  465.                 waterRegion = waterRegions[regionKey]
  466.                 outlet = outlets[regionKey]
  467.                 edge = edges[regionKey]
  468.                 start = findHighestPoint(grid, waterRegion)
  469.                 end = findLowestPoint(grid, outlet)
  470.                 if findHeight(grid, start) > findHeight(grid, end):
  471.                     movedSomething = moveWater(grid, start, end, waterRegion + [end], edge)
  472.                     if grid[start].water == 0:
  473.                         waterRegion.remove(start)
  474.                         if len(waterRegion) == 0:
  475.                             waterRegions.pop(regionKey)
  476.  
  477.         t += 1
  478.  
  479. if __name__ == '__main__':
  480.     import doctest
  481.     # Doesn't test inner functions!
  482.     failures, passes = doctest.testmod(raise_on_error=False)
  483.     if failures == 0:
  484.         main(grid)
  485.     else:
  486.         print 'Stopping with %d failures'%failures
Success #stdin #stdout 0.23s 13456KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
 x #   000000
#  # # 000000
#  # # 000000
#  # # 000000
#    # 000000
###### 000000

 x #   000000
#| # # 050000
#  # # 000000
#  # # 000000
#    # 000000
###### 000000

 x #   000000
#  # # 000000
#N # # 0a0000
#  # # 000000
#    # 000000
###### 000000

 x #   000000
#| # # 050000
#  # # 000000
#N # # 0a0000
#    # 000000
###### 000000

 x #   000000
#  # # 000000
#N # # 0a0000
#  # # 000000
#N   # 0a0000
###### 000000

 x #   000000
#| # # 050000
#  # # 000000
#N # # 0a0000
#~~  # 055000
###### 000000

 x #   000000
#  # # 000000
#N # # 0a0000
#  # # 000000
#~~~ # 066600
###### 000000

 x #   000000
#| # # 050000
#  # # 000000
#N # # 0a0000
#~~~~# 055550
###### 000000

 x #   000000
#  # # 000000
#N # # 0a0000
#  # # 000000
#~~~~# 077770
###### 000000

 x #   000000
#| # # 050000
#  # # 000000
#N # # 0a0000
#~~~~# 077770
###### 000000

 x #   000000
#  # # 000000
#N # # 0a0000
#  # # 000000
#NNNN# 0aaaa0
###### 000000

 x #   000000
#| # # 050000
#  # # 000000
#N # # 0a0000
#NNNN# 0aaaa0
###### 000000

 x #   000000
#  # # 000000
#N # # 0a0000
#~~#~# 033030
#NNNN# 0aaaa0
###### 000000

 x #   000000
#| # # 050000
#  # # 000000
#~~#~# 066060
#NNNN# 0aaaa0
###### 000000

 x #   000000
#  # # 000000
#N # # 0a0000
#~~#~# 066060
#NNNN# 0aaaa0
###### 000000

 x #   000000
#| # # 050000
#  # # 000000
#NN#N# 0aa0a0
#NNNN# 0aaaa0
###### 000000

 x #   000000
#  # # 000000
#N # # 0a0000
#NN#N# 0aa0a0
#NNNN# 0aaaa0
###### 000000

 x #   000000
#| # # 050000
#~~#~# 033030
#NN#N# 0aa0a0
#NNNN# 0aaaa0
###### 000000

 x #   000000
#  # # 000000
#~~#~# 066060
#NN#N# 0aa0a0
#NNNN# 0aaaa0
###### 000000

 x #   000000
#| # # 050000
#~~#~# 066060
#NN#N# 0aa0a0
#NNNN# 0aaaa0
###### 000000
https://ideone.com/JfPbjL
language:
Python (cpython 2.7.16)
created:
8 years ago
visibility:
secret

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!