fork download
  1. from math import*
  2. R=range
  3. L=len
  4. p={'A':(k:=[3,2,1],1),'H':([5]*3,1),'B':(k,0)}
  5. g=lambda x,i:L(set([j[i]for j in x]))
  6. def u(b,x,y,l,r,o,Q,i=0,H=0):
  7. m=lambda X,Y:max(abs(X-x),abs(Y-y))
  8. if i*~-((j:=b[x][y])in p)and int(j)<l[(P:=m(*o[0])-1)]:b[x][y]=l[P]
  9. if(j=='B')*i*(m(*o[0])<4):Q.append((x,y))
  10. for q,w in[(0,1),(1,0),(0,-1),(-1,0)]:
  11. if(X:=x+q)>=((Z:=(X,Y:=y+w))in o)==0<=Y and(r or 1 in[g(O:=o+[Z],0),g(O,1)]):
  12. try:
  13. if(K:=dist(o[0],Z))>=H:u(b,X,Y,l,r,o+[Z],Q,1,K)
  14. except:1
  15. def f(b):
  16. Q,S=[],[]
  17. for x in R(L(b)):
  18. for y in R(L(b[0])):
  19. if(E:=b[x][y])in['A','H']:u(b,x,y,*p[E],[(x,y)],Q)
  20. while Q:
  21. S+=Q;q=[]
  22. for i in Q:u(b,*i,*p['B'],[i],q)
  23. Q=[j for j in set(q)-set(S)]
  24.  
  25. tests = [["""H""", 'H'],
  26. ["""000000000000000
  27. A00A0000000000A
  28. 000000000000000
  29. 000000000000000
  30. 000000000000000
  31. 000000000000000
  32. 000000000000000
  33. 000000000000000
  34. 0000000A0000000
  35. 000000000000000
  36. 000000000000000
  37. 000000000000000
  38. 000000000000000
  39. 000000000000000
  40. A000000000000AA""","""333332100001233
  41. A33A3210000123A
  42. 333332100001233
  43. 222222100001222
  44. 111111100001111
  45. 000011111110000
  46. 000012222210000
  47. 000012333210000
  48. 0000123A3210000
  49. 000012333210000
  50. 000012222210000
  51. 111111111111111
  52. 222100000012222
  53. 332100000012333
  54. A321000000123AA"""],
  55. ["""000000000000000000
  56. 0AAAA00000000AAAA0
  57. 000000000000000000
  58. 000000000000000000
  59. 000000000000000000
  60. 000000000000000000
  61. AA0000000000000000
  62. 000000000000000000
  63. 000000000000000000
  64. 000000000000000000
  65. 00000000000000A000
  66. 000000000000000000
  67. 000000000000000000
  68. 00000A000000000000
  69. 000000000000000000
  70. 000000000000000000
  71. 00000000000000000A
  72. A00A000A0000A00000""","""333333210012333333
  73. 3AAAA32100123AAAA3
  74. 333333210012333333
  75. 222222210012222222
  76. 222211110011111111
  77. 333210000000000000
  78. AA3210000000000000
  79. 333210000001111111
  80. 222210000001222221
  81. 111110000001233321
  82. 00111111100123A321
  83. 001222221001233321
  84. 001233321001222221
  85. 00123A321001111111
  86. 111233321111111222
  87. 222222222222222233
  88. 33333233322333223A
  89. A33A323A3223A32233"""],
  90. ["""H00000000000000000
  91. 000000000000000H00
  92. 000000000000000000
  93. 00000000H000000000
  94. 000000000000000000
  95. 000000000000000000
  96. 000000000000000000
  97. 000000HH0000000000
  98. 000000000000000000
  99. 000000000000000000
  100. 000000000000000H00
  101. 000000000000000000
  102. 000000000000000000
  103. 000000000000000000
  104. HHH00000000H000000
  105. 000000000000000000
  106. 000000000000000000
  107. 000000000000000000""","""H55505555555555555
  108. 555505555555555H55
  109. 555505555555555555
  110. 55550555H555555555
  111. 000555555555555555
  112. 000555555555000000
  113. 000555555555000000
  114. 000555HH5550555555
  115. 000555555550555555
  116. 000555555550555555
  117. 000555555550555H55
  118. 555555005555555555
  119. 555555005555555555
  120. 555555005555555555
  121. HHH55500555H555000
  122. 555555005555555000
  123. 555555005555555000
  124. 555555005555555000"""],
  125. ["""000000000000000000000A
  126. 00000B0000000000000000
  127. 0000000000000000000000
  128. 0B0000000000000000000B
  129. 000000000000000000B000
  130. 0000000000000000000000
  131. 000000000000000000000B
  132. 0000000000000000000000
  133. 000000B00000000A000000
  134. 000000000000000000000B
  135. 0000000000000000000000
  136. 000000B000000000000000
  137. 000000000000000000000B
  138. 0000000000000000000000
  139. 000000000H000000000000
  140. 000000000000000000000B
  141. 0000000000000000000000
  142. A000000000000000000000
  143. A00000000000B000000000
  144. A000000000000000000000
  145. HHHH00000000000000A00H
  146. 00000000A0000000000000""","""000000000000000000123A
  147. 00000B0000000000001233
  148. 0000000000000000001223
  149. 0B0000000000000000123B
  150. 000000000000000000B003
  151. 0000001000001111111003
  152. 000000200000122222123B
  153. 0000003000001233321003
  154. 000123B32100123A321003
  155. 000000300000123332123B
  156. 0000003000001222221003
  157. 000123B555555111111003
  158. 000000555555500000123B
  159. 0000005555555000000003
  160. 111100555H555000000003
  161. 222100555555500000123B
  162. 3321005555555000000003
  163. A555555555555001115555
  164. A55555511111B001225555
  165. A555555222210001235555
  166. HHHH55533321000123A55H
  167. 55555553A3210001235555"""],["""0000000000
  168. 0A00000000
  169. 0000000000
  170. 000B00B000
  171. 0000000000
  172. 000B00B00B
  173. 0000000000
  174. 0B00B00000
  175. 0000000000
  176. 0000000000""","""3332101000
  177. 3A32102000
  178. 3333103001
  179. 223B33B322
  180. 1113103003
  181. 123B33B33B
  182. 0003003003
  183. 0B02B02002
  184. 0001001001
  185. 0000000000"""]]
  186. def to_board(d):
  187. return [list(i)for i in filter(None,d.split('\n'))]
  188.  
  189. for a, b in tests:
  190. board = to_board(a)
  191. f(board)
  192. assert (r1:='\n'.join(''.join(map(str,i))for i in board)) == b
Success #stdin #stdout 0.15s 14244KB
stdin
Standard input is empty
stdout
Standard output is empty