regolith reservoir

approach

the rock paths are expanded into individual cells and stored in a dict keyed by (r, c), which later holds sand too. the world remembers abyss, one row below the lowest rock. drop moves a grain down, down-left or down-right until none of the three cells is free or it reaches the abyss row, then stores it where it stopped and returns whether it rested above that row.

part 1 drops grains until one falls through to the abyss row. part 2 uses the same drop; grains that reach that row just stay there, which acts as the floor, and it counts until a grain comes to rest on the source itself.

the trick

stopping grains at the abyss row instead of losing them gives part 2's floor for free.

    def drop(self, p):
        r, c = p
        while r + 1 <= self.abyss and any((r + 1, c + dc) not in self.particles for dc in range(-1, 2)):
            if (r + 1, c) not in self.particles:
                r += 1
            elif (r + 1, c - 1) not in self.particles:
                r += 1
                c -= 1
            else:
                r += 1
                c += 1