day 17 wip

day17.py

@@ -0,0 +1,110 @@
+from typing import List
+from typing_extensions import Self
+import numpy as np
+from dataclasses import dataclass
+from numpy.typing import NDArray
+from itertools import cycle
+
+TEST = ">>><<><>><<<>><>>><<<>>><<<><<<>><>><<>>"
+INPUT = ">>><<<>>>><<<>>>><<<<>>><<<>><<<<><<<<>><<<<><<<><>>>><<<>>><<<<>>><>>>><<<>>>><>>><>>><<<<>>>><<>>><>>><>>>><<<<>>>><<<>>><>><<<>>><<<>>>><<<>>><>>>><<>>>><<<<>>><<>>>><<<>>>><>>><<<>><<>>>><<>>>><><<<>><<<>>>><><>>>><<><<<<>>>><<<>>>><<<><>>><<<>>>><<<>>><<<>><<<<>><<<<><<<>><<><>>><<<<>>>><<<<>><<>>>><<<<>><<<<>><<<<><<<>><<<>><<<>><<>><<>>><<<<>>><<<<>><<<><<<>><>>>><<<<>>><><<<<>><<<<><<><<<>>>><<<<>><>><<<<>>><<<>>>><<<>>><<<><<><<<>>>><<<<><><<<>>><<<<>>>><<<<>>><<<>>>><<<<><<<>>><<>><<<><>>>><<<<>>>><<<>><<<>>><<<><<<><<<<>>>><<>>>><<<>>>><><>>>><<>>>><<<>><>>>><<<<>>><<<><>>><<<><<<<>>><<<><<<>>><>>><<<>>>><>>><<<>><><>><<<<>><>>><>>>><<>><<<>><<>>><<>>>><<<<>>>><<<<>>>><<<>>><<>>>><<><<>>><<>>><<<<>>><<>><>>>><><><<<<>>>><<<<>><<<<><<<<>>>><<<>>><<><<<>>>><<><><<>>><<>><<<>>><>><<<>><<<<>>><><<>><<<>>><<<<>><>>><<>><>>><<<><>><<>>>><>>><<>>><<<<>><<<<>>>><>>><<<><<<<>>><<>>><>>><<<<>>>><<<<>><<<>>>><<>>><>>><>>><<<>><<<<>>><<<<>>>><>>><>>>><<<<><<><<<<>>>><>><>>>><<<><<<>>><>>><<<<><>>><<<>><<<>>><<<>>>><>>><>>>><<>><<>>><<<<>>>><><<<<>><<<>><<>>><>><><<>><<>>>><<<<>>>><<<>>><<>><>>><<<<><<<<>><<>>>><<<<>>>><<<>><><<<<>>>><><<<><<<<>>>><>>>><<<>>><<<<>>>><<<>>>><<>>>><<><<>><<<<>><<>>>><<<><<>><<<<>>><<<>>><<>>><<<>>>><<>>>><<>>><<<<>>><<<>>>><<>>>><<<<>>><<<>><<<>><<><<<<>><>><>><<<><<>>>><<<<><>>>><<<<><<>>><<>>>><<>>><>>><><<><<<<>>>><<>>>><<<>><<>>>><<<>>><<<>>><<<><<<<>><<<>><<<<><<>>><<<>>><<<>><<<>>>><<><>><<<><<<<>><<>><<<>><<<<>><<<>>>><>><<>>>><<<><<>>>><<<>>>><>>><<<>>>><<<<>>><<<><<>>><<<>>><<>><<>>>><<<<>>><<<><<<><>><>>><<<<><<<<><<<<>><<<<>>><<<<><<><<>>><<<<>><<<>><><>>>><>><<<<>><<>>>><><>>><><<>>><>><<<<>><>>><>><<<<><<<>>><>><<<>>>><><>>><<<>>>><<><><<<<>>><<>><<<>><<><<<<>>>><<>><>>><<>>><<>>>><<<>>><<>>>><<<<>>>><<<<>><<<<>>>><<<<>>><<<><<<>>><>>>><>>>><<<><<<<>>><><<<<>>>><>><<>>>><<<<>>><<>>><><<<<>>><<<<>>>><<<<>>>><<><>>>><<><<>>><<><<<<>>>><<<<>><<<><>>>><>>>><<>>><<<<><><<>>><<>>><><><><>>><<><<>>><>>><<<<>>>><><<<>>>><>>><<<>>><<<>>>><<<>>>><<>>><>>>><<>><<<>>><<<>><<<<><<>>>><<<<>>><<<<>>><>>><<<<>><<<<>><<<<>><<<><<<>>>><<<>>>><>>><<>><<<<><<<>>>><<<>>><<<<>>>><>><<<><<>>>><<>>><<<<>>><>>><<<>><<<<>>>><<<>><<<>>>><>>>><<<><<<<>>>><>>>><<<<>>>><<<<><<>>>><<>>><>>>><<<<>>><>>><>><<<>>><>>>><<<<>><<>>><<<>>>><<>>><<><<<<><<<>><<><<>><<<><<><<<<>><<<<><<><<>>><<<>><<>>>><<>>><<>>>><<<<>>><<<>>><<>>>><<>>>><<<<>>><><>>>><<>>>><<>><<>>><<<>><<<>><>>><<<>><<<>><<<>>><<<<><<<<>>><<<<>>>><<<>>><<>><<>>><<<><<><<>><<<<>><<><<<<>><<<<>>>><>>>><><<<>>>><<<<>>><>><<<>>>><<>>>><>>><<<>><>>>><>><>>>><<<<>>><<><><<<>>><<<<>>>><<<><>><>><<<>>>><>>>><>>>><><<<<>><<<<><><<<>>>><<>>><<<<>>><<<>>>><>>><<><<>>>><<>><<>><<<<>>>><<>>>><><<<<><<><<<<>><<>><<<<><<<><<<<>>>><<<>><<<<>><<<<>>>><<>>><<<>><>>><>>><<>><>>>><><<>><<<><<>>>><<>>><<>>><<>>><<<<>>>><<>>><>><>>><<<><<<>><<>>>><<<<>>>><<<<>>>><<>>><<<>>>><<<<>><<<<>>>><>>>><>>><>>>><<<<>><><<><<<<>>><<<>>><<>>>><>><><<<>>>><<<<>>>><<>><<>><<<><<<<><<<<>><<<<>>>><>>>><>>>><<>>><<<>>><<<>>>><<>><<<>><<<><>>><<>>>><<<><>>>><><<<<><<>><><>><>><<<>><<>>>><<>>><<<><<>>>><<<<>>>><>>>><<<<>>><<<<>>><<><<<<>><>>>><<>>>><>>>><<<>>><>><<<>><<><<<<>><<<>>>><<<<>><<<<>>><<<>>><<<>>>><<<<>>><<<>><<<<>>><<<>>>><>>>><<><>><<>>>><<<<>>><><<<<>>><><>>>><>><<<<>>><<<<>>>><<><<<>>><>>>><<<>><<>>>><<<<>><><<>>><<>>>><>><<<><<>>><<<>>><<<<>><><<><><<<<>><<<<>><<>>><<<>>>><<<<>><<<>><<<>>><<><<<<>>><<<>>><>><>>><>><>><<<<>>>><>><<<<>><<>>>><>>>><>>>><<<<>>><<<<><<>>><>>><<<<>>>><<<>>>><<>>>><<<><<<><<<<><>>>><<<>><<<>>>><<<<>><<>>><<<>>>><<<><<<>>><<><><<>>><<<>><<<<><<<<>>><<><<<>><><>>>><<>><<<<><<<<>><<<><<><<<<>>>><<>>>><<<<>><<<>>><<><>>>><<>><>><<<<>><<<>>>><>><<<<>>><><<<><<>>><>>>><<<<>>><<<><>>><>>>><<<<>>><<>><<>>><<<<><<<<>><<><<><<>><>>><<<>>>><>>><>>>><<<<><<>><>>><<<>>><>>>><<<<>>>><<>>><>>><<<<>>>><<>>>><<<>>><<<>>>><<<>><>>><<>>><<<><<<<>>>><<<>>><<>><<<<>><>>>><<>>><<<<>><<<>>>><<>>>><<>><><<<<>>><<<<>><<<>><<<<>>>><<>>><<><<>>><<>><>><>>>><<<<><<>>><<<<><<<<>><<><<>>><<<>>><<<<>><<<<>>><>><<<<>>><<<>>>><>>><>><<<>>><<>>>><<<<>>>><<>>><<<<>><<><>>><<>>><<<<>>><<<<><<>><<<<><<<>>>><<>>>><>><<<><<<<>>>><<<<>><<<<><<>>>><<<>>>><>>><>><<<<><<<<>>>><<<>>><<<>>>><><>>><<<<><<<<>>>><<<>>><<>>>><<><>><<><<<<>>>><>>>><<<>><<<>><<<<>>><<<<>><<<>><>>><<<>>>><<<<>><<<<>>>><<>>><<<>>>><<<<>>>><<<>><>><<>><><<<<>><<><><<>><<>><<<<>>><<>>>><<<<>>>><<><<<<>>>><<<>><<<<>>>><<<>>><<><>>>><<<>>><>>>><<<<>>>><>><<><<<<>>><<>>><<<<>><<<>>>><<>>>><<<><<<>>>><<<<><<<>><<<<>>><<<>><<<><<<>><>>><<<<>>><<<>>><<<>><<<>>>><<<<>>><<<<>>>><<<<>><<<<><<<>>><><<<>><<<>><<<>><<<>>><<<><<<>>>><<<>><<<<>>>><><<<<>><<<<>>><<<<>><<><<<>>>><>>><><<<<>>><>>><<<<>>>><<<<>>><<<>>>><><<>><><<><<<<><<<<>>><>><>>>><<<<><<<<>><<<><>>>><<>>><<><<>>><<>>>><<<<>>>><<><>>>><<<<><<>><<><<>><<<<>>><<<<>><<<<>>>><>><<><<<>>>><<<>><<>>>><<<><<<<>>>><<<>>><<>><<<>>>><<>><<<<>>>><<<<>><<<>>>><>><<<<><<<>>>><<<><<<>><>>><<>>>><<<<>>><<<><<><>><<><>><<<><><<<><<<<>>><<>><<<><<>><<<><<<>><<<>>><>>>><<<>>>><<<<>>>><<<>><>><<>>><<<><>>><<<>><<<><<>><>>>><><>>>><<<>>><<<>>>><>>><<<<>>><>>>><<<>>>><<>>><<>>>><<>><><>>>><<<>>>><>>><>><<>>><<>>>><>><<>>><<><>>>><<<>><<>>><><<<>><<>>>><<<<>>>><>>><<<<>><>>><<>><<><<<>>><>><<<><<>><<>>><<><><<<<>>>><<<>>><>>>><>><<<>>><<<>>><<<<>>>><><<<<>>><<<<><<<>><<<<>>>><>>>><<>><<<<>>>><<>>><>><<><<<>>><<<>>>><<<<>>><<<>>><>>><<<>>><<<<>><<<><<>>>><><<><<>><>><<<<>>>><>><<>><<<<>>><<>>>><<<>><<<>>><<<>>><<<<>>><<<<>>>><<<<>><>>>><<<<>>>><>>>><>>><<<<><<<<>>><><<>>><<>>>><<<<>><<<<>>><<<>><<<>>>><<<<>>><<<>>>><<<><>>>><>>><<>>>><<<><<<><>>>><<<>>>><<>>>><>>>><<><<<<>><<<<><>>>><<><<<<>>>><>>><<<<>>><<>>>><<<><<<<><><<>>><<<>>><<>><<<>>>><>>>><<<<>>>><<<>>><<<>><<><<<>>>><<<<>><<<>>><>>><>><<<>>>><<<<>>>><<><<<>>>><>>>><<<>>>><<>><<<<>>><<<>>><><><<<<>>>><<<>>>><><<>><<>>>><<<<>>>><<<><<>>>><>>><<<<>><<>><<<>><><<<<>>><><<<>>>><<<><<<>>><<<<>><<<<>>>><<<>>>><<>><<<>>><<<<>>><><<>><<<<>>>><<<<>>><<>>><<<<><>><>>><<<<><>>>><<>><>><<<>>>><<<>>><<>><<><>>>><<<>>>><<<<>>>><<>>>><<<<>>>><<<>>>><<<>><<<>>><>>><>>><>>>><<>>>><<<>>><<<<>>><<>><<<><<<<>><>>><<<>>><<<>>><><<>>><>>>><<<<>>><<>>>><<<>>>><<<<>>><<<<>>>><<<>>><<><<>>><<>>>><<<>>><<><<<>>>><><<<<><><<><<>><>>>><<>>><><<><<<<>>><<<>>><<<<>>>><>><<<<>><<<<>><<><<<>>><>>>><>>>><<<<>>>><<<<><><<<><>>><<<<>><>><><<<<>>><<<><<><<<<>><<><>>>><>>>><<<>>>><<><<>>><<>><<<<>>><<<<>>>><>>>><<<<>><<>>><<<>>>><<<>>>><<<>>>><<<<>>><<<><<<><>>><<<><>>><<<><>>><>><<<>>><<<>>><<<<>>>><<>><<<>><>><<<<>>><<<>>>><<><<<>>>><<<<><><>>>><<>><>>>><<>>><<>>>><<>>>><>>><>>><<<>>><<>>><>><>>>><<<>>><>><>>>><<<>>><<><<<>>>><<<>><<><<<>>><<<<>>><<<<><<<<>>>><<<>><<<<>>><<<<>>><<>><<<<><<<><<<<>>>><<<<><>>>><>>>><<><<<<>><<>><<>>><>>><>>>><<<>>>><>><<<<>>>><<<<>><<<<>><<<>><<<<><>>>><>><<<>>><>><>>><>>><<<<>>><<>><<<>>><<<<>>><<><<<>>><<<>>>><<>>><<<<>><<<<>>><<><<>>><<<<>>><<<>>><<<<>>><<>>><<<>><<<>>>><<<>><<<<>>>><<>><>>>><<>>><>><<<>><<<<>>>><<>><<<<><<<>>>><<<>>><<>><<<>>>><<>>><<<<>><<>><<>><<<<>>>><<<>>><<<<>>>><<<>>><>>>><<<>>>><<>><<<<>><<>>>><<<<><<<>><<<<><<>>>><<<<>>>><<>>><<<<><<<>>>><<<>><<><<<<><><<>>>><><<<<><<<><<<<>><<><<<<><<<<>><<<>><<>><<<<><<>>><<><<<>><<>>><<<>>>><<>><<>><<<<>>><<<<>><<>>>><<<>>><<<<>><<>>><<><><<><<>>>><>><>>><<<><<<<>><<><<>>><<<<><<>>><>>>><<<><<><<>>>><><>>>><<>><<><>>><<><<<<>><<>>>><<>>><<<>>>><<<><>>><<>>>><>><<><<>><<<><<<>><<>><>><><><<>>>><<<>>><>>><<<<>>><<<>>><<<>><>>><<<<>><><<<><>>>><<>>><>>><<<>>><<<>><>>>><<>>><<<><<<>>>><><<<<>>><<<<>><<>><<<><<<<>><<<>>>><<<><>><<<<><<<<>>><<<>>>><<<<><<<>>><>>>><<<<>>><<><<>>>><<><<>><<>><>><<>>>><<<<>>><>>>><<>>>><<><<<>>><>>><<<>>>><<<<>><<<<><<<><<<><<<><<<>>>><<><<><><>><>>><<<>><>>><<<><>>><<>>><<<<><<<>>><<<<>>><<<<>>><<>>>><<<>>>><<<>><<<<><<<<>><<<><<<><<<>><<>>>><<<>><<<><>>><>>><<><<<>>><<<<>><<<>>>><>><<<>><>>><><<<><<<<><<<><<<>>><<><<>>>><<>>><><<<<>>>><<<>>>><<>><<<>><<>>><<><<<>><<>><<>><<<<>><<<<>><<<>>>><>>>><<>>><<<<><<>><<<>><<>>>><<<<>>><>><>><<<>>>><>>><<<>>><<>><<<<>><<><<<><<<<>><<<<><>><<<>>>><<<>>>><<>>><>><<>>>><<<>><<<>>>><><<<><<>><>><<>><<>>>><<>>><<<<><<><>>><><<<>>><<>>><<<>>><<<>>>><>>><<<<>>>><<<>>><<>><<<>>><<<>>><>>>><<><<<<>>>><<<><<<<>>><<<<>>><>><<<>>>><<><<><<<>>>><><<<>><<<<><<<>>><>>><>><<><<>><<<<>>>><>>><><<<<>>><<>><<<>>><><<<>>>><<<<>>><<>><<<><<><<<<>>>><<<>>>><<<>>>><><>><<<>>>><<><>>>><<<<>>>><<<>>>><<>>><<<<>>>><>>><>>><<>>>><<<><<>>>><><<<<>><<<>>>><<><<<<>>>><<>>>><><<<>>>><>><<<>>>><<>>><<>>><><<<>>><<<>><><<<<>>><<>>>><>>><<<>>><<>>>><<>><<<<><<>><<<>><<<<>>><>><<<<>>><<<>>>><>>>><<<>>>><>>><<<<>>><>><<<<><<<>><>>><>>>><<>>>><<>><<<>><<>>><<<>><<<<><<>><<<<>><<<<>><<<><>>>><>><><<<><<><>>><<<<><<<>><<>>><>>><<<<>><<<><<<>><>>>><<<<><><<>><<<<>><<<<>><>>>><><<<><>>><<<<>>>><<<>>><<>>>><<<>><<<>><>><<<><<<<>>><<<>>>><<<<>>>><<<>><<<<>>><>><<<<>><><<<<>><><<<<>><<>><<<<>>><>>>><<<<>>>><>>><>><>>>><<<>>>><><<<><<<>><<<<>>><<<<>><<<>>><<<<><<<>>><<<>>><><>>>><<<><<<<><<<>><<<<>>><<><<<<>><>>><<<<>>><>><>>>><><<<>>>><<><>>>><<><<<>>><<>>><<<<><>>><<><<>>><<<>>><><<<>>><>>><<<<>>>><<<>>>><<>>><<<<>>>><<>><<<<>><<<<><<>><<<<><<<>>>><<<><>>>><<>>>><<<>>>><<<>><<<>>><<>>>><>>><<<<>><<<<>><>>><<<<>><<>>><<<<><<<<>><<<<>>><<>><<><<<>><<>>>><<<<>><<<<>>><>>>><<>><>>>><<<><<<<>>><>>>><<<>><<<<>><<>>><<<>><<><>><<><>>>><>>>><<<<><>>><>><<<<>>>><<<>>>><<<>>>><<<>>>><<<><<<>>><<<>><<<<><<<><<>><<<>>>><<<>>>><<<<><>>><<><<<<>>>><>><<>>>><<>>>><>>><<<>>>><<><<<<>>>><>><><<<<>><>>><>>>><<<>>><<<<>>><<<<><<<>>>><>><<<><>>><<<<>>><>><>>><<<>><>><<><<<>><>>>><<><<<>>>><><<<>>>><<>><<>>>><>><><<<<><<<>>>><<<><<<>><<<<>><<><<>><>>>><<<<>>>><>><<<<>>>><<<<><>>>><<<><<<<><>><<>><<<>>><>><<<<>>><><<<>>><>><<>><<>>>><<<<>>><<<>>>><<<<><>><<<><<<><<><<<><<<<>>>><><><<><<>>><<>>><<>><>>><>>>><<><<<><>>><<<<>><<><<>><<<<>><<>>>><<<>>>><>>><<<>>><<<>>><<<>><>><<<<>>>><<>>>><<<>>>><>><<<>>>><<<>>>><<<>>><<>><<>><<<<><<><>><><<>><<>>>><<<>><>>><><<><<<><<>><<>>><<<>>>><<><>>>><<<<><<>>><><<<<>><<<>><<<<><<><<>>><>><>>>><<>>>><<<>>><<>>>><<<<>>>><<<>>>><>>>><>>><<<><>><<><>><<<<>>><<<<>>>><<<<>>>><>><><<<>>>><<<<>><<>>>><<<<><><<<>>>"
+SHAPES = [
+    np.ones((1,4)),
+    np.array([[0,1,0], [1,1,1], [0,1,0]]),
+    np.array([[0,0,1], [0,0,1], [1,1,1]]),
+    np.ones((4,1)),
+    np.ones((2,2)),
+]
+WIDTH = 7
+
+
+@dataclass
+class Rock:
+    shape: NDArray
+    y: int
+    x: int = 2  # default start two left from the wall
+
+    @property
+    def width(self):
+        return self.shape.shape[1]
+
+    @property
+    def height(self):
+        return self.shape.shape[0]
+
+    def overlaps(self, other: Self):
+        x_min = min(self.x, other.x)
+        y_min = min(self.y, other.y)
+        x_max = max(self.x + self.width, other.x + other.width)
+        y_max = max(self.y + self.height, other.y + other.height)
+        self_ = np.zeros((x_max - x_min, y_max - y_min))
+        other_ = self_.copy()
+        self_[self.x - x_min, self.x + self.width - x_min:self.y - y_min, self.y + self.height - y_min] = self.shape
+        other_[other.x - x_min, other.x + other.width - x_min:other.y - y_min, other.y + other.height - y_min] = other.shape
+        return 2 in self_ + other_
+
+
+def hit(rocks: List[Rock], rock: Rock):
+    return next((value for other in reversed(rocks) if (value := other.overlaps(rock))), False)
+
+
+def push(rocks: List[Rock], rock: Rock, dir, width, height: int):
+    if dir == ">":
+        if rock.x + rock.width + 1 > width:
+            # hit wall
+            return
+        if rock.y > height:
+            rock.x += 1
+            return
+
+        # find whether there's a rock to hit
+        rock.x += 1
+        if hit(rocks, rock):
+            rock.x -= 1
+
+    elif dir == "<":
+        if rock.x - 1 < 0:
+            # hit wall
+            return
+        if rock.y > height:
+            rock.x -= 1
+            return
+
+        # find whether there's a rock to hit
+        rock.x -= 1
+        if hit(rocks, rock):
+            rock.x += 1
+
+
+def fall(rocks: List[Rock], rock: Rock, height: int):
+    if rock.y > height:
+        rock.y -= 1
+        return False
+
+    if rock.y - 1 < 0:
+        return True
+
+    rock.y -= 1
+    if hit(rocks, rock):
+        rock.y += 1
+        return True
+    return False
+
+
+if __name__ == "__main__":
+    next_shape = cycle(SHAPES)
+    height = 0
+
+    rocks = []
+    rock = Rock(next(next_shape), 3)
+    landed = False
+    for dir in TEST:
+        rock = rock if not landed else Rock(next(next_shape), height + 3)
+        print(dir, rock, WIDTH, height)
+        push(rocks, rock, dir, WIDTH, height)
+        landed = fall(rocks, rock, height)
+        print(" ", rock, "" if not landed else "has landed")
+        if landed:
+            rocks.append(rock)
+            height = max(height, rock.y + rock.height)
+
+    print(rocks)