|
| 1 | +import argparse |
| 2 | +import itertools |
| 3 | +import os.path |
| 4 | +import textwrap |
| 5 | +from copy import deepcopy |
| 6 | +from dataclasses import dataclass |
| 7 | +from typing import ClassVar |
| 8 | + |
| 9 | +import pytest |
| 10 | + |
| 11 | +from support import Direction4 |
| 12 | +from support import timing |
| 13 | + |
| 14 | +INPUT_TXT = os.path.join(os.path.dirname(__file__), 'input.txt') |
| 15 | + |
| 16 | +# NOTE: paste test text here |
| 17 | +INPUT_S = '''\ |
| 18 | +>>><<><>><<<>><>>><<<>>><<<><<<>><>><<>> |
| 19 | +''' |
| 20 | +EXPECTED = 3068 |
| 21 | + |
| 22 | +Coord = tuple[int, int] |
| 23 | + |
| 24 | + |
| 25 | +@dataclass |
| 26 | +class Rock: |
| 27 | + SHAPES_STR: ClassVar[str] = textwrap.dedent("""\ |
| 28 | + #### |
| 29 | + |
| 30 | + .#. |
| 31 | + ### |
| 32 | + .#. |
| 33 | + |
| 34 | + ..# |
| 35 | + ..# |
| 36 | + ### |
| 37 | + |
| 38 | + # |
| 39 | + # |
| 40 | + # |
| 41 | + # |
| 42 | + |
| 43 | + ## |
| 44 | + ## |
| 45 | + """) |
| 46 | + name: str |
| 47 | + shape: set[Coord] |
| 48 | + |
| 49 | + @classmethod |
| 50 | + def from_str(cls, name, s): |
| 51 | + """Load the rock coords from string. |
| 52 | +
|
| 53 | + Must load the coords "bottom up". |
| 54 | + """ |
| 55 | + coords = set() |
| 56 | + for y, line in enumerate(reversed(s.splitlines())): |
| 57 | + for x, c in enumerate(line): |
| 58 | + if c == '#': |
| 59 | + coords.add((x, y)) |
| 60 | + return cls(name, coords) |
| 61 | + |
| 62 | + def move(self, _dir: Direction4, n=1): |
| 63 | + self.shape = {_dir.apply(*coord, n=n) for coord in self} |
| 64 | + |
| 65 | + def fall(self): |
| 66 | + """Direction is reversed due to the up-down flip.""" |
| 67 | + self.shape = {Direction4.UP.apply(*coord) for coord in self} |
| 68 | + |
| 69 | + def __iter__(self): |
| 70 | + yield from self.shape |
| 71 | + |
| 72 | + def __str__(self): |
| 73 | + return f'Rock({self.name})' |
| 74 | + |
| 75 | + __repr__ = __str__ |
| 76 | + |
| 77 | + |
| 78 | +class Chamber: |
| 79 | + jet_to_dir = { |
| 80 | + '<': Direction4.LEFT, |
| 81 | + '>': Direction4.RIGHT, |
| 82 | + } |
| 83 | + |
| 84 | + def __init__(self, jets: str, rocks: list[Rock]): |
| 85 | + self.jets = itertools.cycle(jets) |
| 86 | + self.rocks = itertools.cycle(rocks) |
| 87 | + self.occupied: set[Coord] = set() |
| 88 | + self.rock = None |
| 89 | + |
| 90 | + @property |
| 91 | + def max_height(self): |
| 92 | + if not self.occupied: |
| 93 | + return 0 |
| 94 | + return max(c[1] for c in self.occupied) |
| 95 | + |
| 96 | + def collision(self, rock: Rock): |
| 97 | + in_wall = any(c[0] in [0, 8] for c in rock) |
| 98 | + on_the_floor = any(c[1] == 0 for c in rock) |
| 99 | + in_occupied = any(c in self.occupied for c in rock) |
| 100 | + return any((in_wall, on_the_floor, in_occupied)) |
| 101 | + |
| 102 | + def spawn_rock(self, rock): |
| 103 | + rock = deepcopy(rock) |
| 104 | + rock.move(Direction4.DOWN, n=self.max_height + 1 + 3) |
| 105 | + rock.move(Direction4.RIGHT, n=3) |
| 106 | + self.rock = rock |
| 107 | + |
| 108 | + def process_rock(self): |
| 109 | + """Repeat following. |
| 110 | +
|
| 111 | + - spawn rock |
| 112 | + - move |
| 113 | + - fall |
| 114 | + """ |
| 115 | + rock = next(self.rocks) |
| 116 | + self.spawn_rock(rock) |
| 117 | + |
| 118 | + while True: |
| 119 | + _dir = self.jet_to_dir[next(self.jets)] |
| 120 | + |
| 121 | + # move the rock |
| 122 | + self.rock.move(_dir) |
| 123 | + if self.collision(self.rock): |
| 124 | + # if collision, revert |
| 125 | + self.rock.move(_dir.opposite) |
| 126 | + |
| 127 | + # fall the rock |
| 128 | + self.rock.fall() |
| 129 | + if self.collision(self.rock): |
| 130 | + # collision during the fall, freeze and continue |
| 131 | + self.rock.move(Direction4.DOWN) |
| 132 | + self.occupied |= self.rock.shape |
| 133 | + break |
| 134 | + |
| 135 | + def _coord_to_char(self, c): |
| 136 | + if c in self.occupied: |
| 137 | + return '#' |
| 138 | + if self.rock and c in self.rock: |
| 139 | + return '@' |
| 140 | + return '.' |
| 141 | + |
| 142 | + def __str__(self): |
| 143 | + print_height = self.max_height + 8 |
| 144 | + rows = [] |
| 145 | + for y in range(print_height, 0, -1): |
| 146 | + row = str(y) + ' |' + ''.join(self._coord_to_char((x, y)) for x in range(1, 8)) + '|' |
| 147 | + rows.append(row) |
| 148 | + rows.append('0 +-------+') |
| 149 | + return '\n'.join(rows) |
| 150 | + |
| 151 | + |
| 152 | +def compute(s: str) -> int: |
| 153 | + jets: str = s.strip() |
| 154 | + shapes: list[str] = Rock.SHAPES_STR.split('\n\n') |
| 155 | + names: str = '-+jio' |
| 156 | + rocks = [Rock.from_str(n, s) for n, s in zip(names, shapes)] |
| 157 | + chamber = Chamber(jets, rocks) |
| 158 | + |
| 159 | + for _ in range(2022): |
| 160 | + chamber.process_rock() |
| 161 | + # print(chamber) |
| 162 | + |
| 163 | + return chamber.max_height |
| 164 | + |
| 165 | + |
| 166 | +@pytest.mark.solved |
| 167 | +@pytest.mark.parametrize( |
| 168 | + ('input_s', 'expected'), |
| 169 | + ( |
| 170 | + (INPUT_S, EXPECTED), |
| 171 | + ), |
| 172 | +) |
| 173 | +def test(input_s: str, expected: int) -> None: |
| 174 | + print() # newline in test output, helps readability |
| 175 | + assert compute(input_s) == expected |
| 176 | + |
| 177 | + |
| 178 | +def main() -> int: |
| 179 | + parser = argparse.ArgumentParser() |
| 180 | + parser.add_argument('data_file', nargs='?', default=INPUT_TXT) |
| 181 | + args = parser.parse_args() |
| 182 | + |
| 183 | + with open(args.data_file) as f, timing(): |
| 184 | + print(compute(f.read())) |
| 185 | + |
| 186 | + return 0 |
| 187 | + |
| 188 | + |
| 189 | +if __name__ == '__main__': |
| 190 | + raise SystemExit(main()) |
0 commit comments