24_hvac.py 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849
  1. #!/usr/bin/env python3
  2. import sys
  3. from collections import deque
  4. from itertools import chain, permutations
  5. def read_grid(f):
  6. rows = [line.rstrip() for line in f]
  7. return list(chain.from_iterable(rows)), len(rows[0])
  8. def shortest_paths(grid, w):
  9. def bfs(start):
  10. dists = [0] * nlocs
  11. work = deque([(start, 0)])
  12. explored = {start}
  13. while work:
  14. loc, dist = work.popleft()
  15. node = grid[loc]
  16. if node.isdigit():
  17. dists[int(node)] = dist
  18. for nb in (loc - 1, loc + 1, loc - w, loc + w):
  19. if grid[nb] != '#' and nb not in explored:
  20. explored.add(nb)
  21. work.append((nb, dist + 1))
  22. return dists
  23. nlocs = max(int(x) for x in grid if x.isdigit()) + 1
  24. return [bfs(grid.index(str(i))) for i in range(nlocs)]
  25. def visit_steps(dists, back):
  26. dists = shortest_paths(grid, w)
  27. best = 100000000
  28. for path in permutations(range(1, len(dists))):
  29. total = 0
  30. prev = 0
  31. for loc in path:
  32. total += dists[prev][loc]
  33. prev = loc
  34. if back:
  35. total += dists[prev][0]
  36. best = min(best, total)
  37. return best
  38. grid, w = read_grid(sys.stdin)
  39. dists = shortest_paths(grid, w)
  40. print(visit_steps(dists, False))
  41. print(visit_steps(dists, True))