24_hvac.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  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 = prev = 0
  30. for loc in path:
  31. total += dists[prev][loc]
  32. prev = loc
  33. if back:
  34. total += dists[prev][0]
  35. best = min(best, total)
  36. return best
  37. grid, w = read_grid(sys.stdin)
  38. dists = shortest_paths(grid, w)
  39. print(visit_steps(dists, False))
  40. print(visit_steps(dists, True))