|
| 1 | +with open("input") as f: |
| 2 | + ls, moves = f.read().strip().split("\n\n") |
| 3 | + |
| 4 | +dirs = {"^": -1, "v": 1, "<": -1j, ">": 1j} |
| 5 | +moves = list(map(dirs.__getitem__, moves.replace("\n", ""))) |
| 6 | +board = {i + 1j * j: x for i, l in enumerate(ls.split("\n")) for j, x in enumerate(l)} |
| 7 | + |
| 8 | + |
| 9 | +def solve(part2): |
| 10 | + walls = {z for z, x in board.items() if x == "#"} |
| 11 | + boxes = {z for z, x in board.items() if x == "O"} |
| 12 | + robot = next(z for z, x in board.items() if x == "@") |
| 13 | + if part2: |
| 14 | + walls = {w.real + 2 * w.imag * 1j for w in walls} |
| 15 | + walls |= {z + 1j for z in walls} |
| 16 | + lefts = {b.real + 2 * b.imag * 1j for b in boxes} |
| 17 | + rights = {z + 1j for z in lefts} |
| 18 | + boxes = lefts | rights |
| 19 | + robot = robot.real + 2 * robot.imag * 1j |
| 20 | + for dz in moves: |
| 21 | + to_move = set() |
| 22 | + to_check = [robot + dz] |
| 23 | + while to_check: |
| 24 | + z = to_check.pop() |
| 25 | + if z in boxes: |
| 26 | + to_move.add(z) |
| 27 | + to_check.append(z + dz) |
| 28 | + if part2 and dz.real: |
| 29 | + other = z + 1j if z in lefts else z - 1j |
| 30 | + to_move.add(other) |
| 31 | + to_check.append(other + dz) |
| 32 | + elif z in walls: |
| 33 | + break |
| 34 | + else: |
| 35 | + boxes -= to_move |
| 36 | + boxes |= {w + dz for w in to_move} |
| 37 | + if part2: |
| 38 | + to_move_lefts = to_move & lefts |
| 39 | + lefts -= to_move_lefts |
| 40 | + lefts |= {w + dz for w in to_move_lefts} |
| 41 | + robot += dz |
| 42 | + return sum(z.real * 100 + z.imag for z in (lefts if part2 else boxes)) |
| 43 | + |
| 44 | + |
| 45 | +# Part 1 |
| 46 | +print(solve(False)) |
| 47 | + |
| 48 | +# Part 2 |
| 49 | +print(solve(True)) |
0 commit comments