Zadanie

W tym zadaniu znowu mamy do czynienia z robotem, teraz jednak przesuwa on się po planszy i jeśli natrafi na pudełko to przesuwa je o jedno pole w kierunku swojego ruchu razem ze wszystkimi pudełkami które są za nim (jeśli pudełek nie blokuje ściana).

Wejście zawiera planszę i listę kroków, gdzie # to ściana, O to pudełko, @ to robot:

########
#..O.O.#
##@.O..#
#...O..#
#.#.O..#
#...O..#
#......#
########

<^^>>>vv<v>>v<<

Rozwiązanie

Parsowanie

Jako że w naszym wejściu jest więcej pustych pól niż zajętych, uznałem że lepiej będzie przechowywać koordynaty pudełek i ścian w słowniku, zamiast zapisywać całego 2d grida. Uznałem też że znowu użyję tricka z trzymaniem koordynatów i wektorów przemieszenia w obiekcie complex, żeby kod był nieco zwięźlejszy.

WALL = 0
BOX = 1
DIR_MAP = {"<": complex(0, -1), "^": complex(-1, 0), ">": 1j, "v": 1}

def parse(data: str):
    grid2d, dirs = data.split("\n\n")
    robot = -1

    grid = dict()
    for i, row in enumerate(grid2d.split("\n")):
        for j, c in enumerate(row):
            comp = complex(i, j * 2)
            if c == "#":
                grid[comp] = WALL
                grid[comp + 1j] = WALL
            elif c == "O":
                grid[comp] = BOX
            elif c == "@":
                robot = comp

    dirs = [DIR_MAP[d] for d in dirs.replace("\n", "")]

    assert robot != -1
    return grid, dirs, robot

if __name__ == "__main__":
    data = sys.stdin.read()
    grid, dirs, robot = parse(data)

Przesuwanie

Kiedy mamy tak ładnie sparsowany input to przemieszczanie robota nie jest trudne, jedynie przemieszczanie pudełek jest nieco bardziej skomplikowane, zauważmy że nie trzeba przemieszczać każdego pudełka w linii, można po prostu usunąć pudełko z początku linii i ustawić inne pudełko na końcu linii.

for d in dirs:
    new = robot + d
    if new not in grid:  # jeśli pole jest puste
        robot = new
        continue

    cur = new
    while cur in grid and grid[cur] == BOX: # iterujemy aż dojdziemy do końca pudełek
        cur += d

    if cur not in grid: # jeśli za pudełkiem jest puste pole
        robot = new
        del grid[new]   # usuwamy pudełko z początku linii
        grid[cur] = BOX # ustawiamy pudełko na końcu linii

    # else na końcu linii jest ściana -> nie robimy nic

Liczenie wyniku

Na końcu musimy zliczyć sumę koordynatów GPS naszych pudełek, które są zdefiniowane jako 100 razy koordynat y plus koordynat x.

result = 0
for c in grid:
    if grid[c] == BOX:
        result += c.real * 100 + c.imag

print(result)

Łącząc wszystko w całość otrzymujemy poprawne rozwiązanie.

Część druga

W drugiej części sprawa się dość mocno komplikuje, tym razem oś y jest rozciągnięta, każdy fragment ściany i każde pudełko zajmuje teraz dwa pola, robot zostaje tylko jeden.

Tym razem jedno pudełko może przesunąć kilka innych pudełek:

###########
##.......##
##..[][].##
##.[][]..##
##..[]...##
##..@....##
###########

Jeśli robot przesunie się do góry, to wszystkie pudełka przesuną się do góry, jednak jeśli przynajmniej jedno z nich będzie blokowane, to nie przesunie się żadne z nich.

Pomysł

Moim pomysłem jest zastosowanie tutaj rekurencji, nie widzę innej opcji gdyż tak na prawdę, nasze pudełka mogą utworzyć łańcuch o praktycznie nieograniczonej długości. Rozdzielimy problem na dwie funkcje:

funkcja sprawdzająca czy możemy przesunąć pudełka
funkcja przesuwająca pudełka (tylko jeśli 1 funkcja zwróciła True)

Lekkim utrudnieniem będzie to że zapisuję pozycje pudełek jako pojedyńczy koordynat, trzeba będzie wziąć pod uwagę to że zajmują one teraz po dwa pola.

Widzę że jest tu dużo pola na błędy więć zastosuję metodologię TDD i napiszę sobie kilka testów jednostkowych dla tych funkcji.

import pytest

from part2 import BOX, WALL, can_box_move, move_boxes, parse

EXAMPLE = """#######
#...#.#
#.....#
#..OO@#
#..O..#
#.....#
#######

<vv<<^^<<^^
"""


"""
VISUALIED STARTING STATE:
##############
##......##..##
##..........##
##....[][]@.##
##....[]....##
##..........##
##############
"""

GRID, DIRS, ROBOT = parse(EXAMPLE)


def test_can_box_move():
    assert can_box_move(complex(3, 9), -1j, GRID)
    assert can_box_move(complex(3, 8), -1j, GRID)

    assert can_box_move(complex(3, 7), 1, GRID)
    assert can_box_move(complex(3, 6), 1, GRID)

    assert can_box_move(complex(3, 7), 1j, GRID)
    assert can_box_move(complex(3, 6), 1j, GRID)

    assert can_box_move(complex(4, 7), -1, GRID)
    assert can_box_move(complex(4, 6), -1, GRID)


def test_move_boxes_left():
    grid = GRID.copy()
    assert 3 + 9j not in grid
    assert grid[3 + 8j] == BOX

    move_boxes(3 + 9j, -1j, grid)
    assert 3 + 9j not in grid
    assert 3 + 8j not in grid

    assert grid[3 + 7j] == BOX
    assert 3 + 6j not in grid

    assert grid[3 + 5j] == BOX


def test_move_boxes_right():
    grid = GRID.copy()
    assert grid[3 + 6j] == BOX
    assert 3 + 7j not in grid
    move_boxes(3 + 6j, 1j, grid)

    assert 3 + 6j not in grid
    assert grid[3 + 7j] == BOX

    assert 3 + 8j not in grid

    assert grid[3 + 9j] == BOX
    assert 3 + 10j not in grid


@pytest.mark.parametrize("x", [0, 1j])
def test_move_boxes_up(x):
    grid = GRID.copy()
    assert grid[4 + 6j] == BOX
    assert 4 + 7j not in grid

    move_boxes(4 + 6j + x, -1, grid)

    assert 4 + 6j not in grid
    assert 4 + 7j not in grid

    assert grid[3 + 6j] == BOX
    assert 3 + 7j not in grid

    assert grid[2 + 6j] == BOX
    assert 2 + 7j not in grid


@pytest.mark.parametrize("x", [0, 1j])
def test_can_bo_move_into_wall(x):
    box = 1 + 1j
    grid = {
        1 + 0j: WALL,
        box: BOX,
        1j: WALL,
        2 + 2j: WALL,
        1 + 3j: WALL,
    }
    assert not can_box_move(box + x, -1j, grid)
    assert not can_box_move(box + x, 1j, grid)
    assert not can_box_move(box + x, 1, grid)
    assert not can_box_move(box + x, -1, grid)

Implementacja

Napiszę kilka funkcji pomocniczych.

def empty(c, grid):  # sprawdź czy pole jest puste
    return c not in grid and not (c - 1j in grid and grid[c - 1j] == BOX)


def is_wall(c, grid): # sprawdź czy pole jest ścianą
    return c in grid and grid[c] == WALL


def get_box(c, grid): # zwróc koordynat lewej części pudełka (lub None jeśli nie ma)
    if c in grid and grid[c] == BOX:
        return c
    if c - 1j in grid and grid[c - 1j] == BOX:
        return c - 1j
    return None

Po jakimś czasie udało mi się zaimplementować funkcję can_box_move, która przecodzi testy jednostkowe które napisałem:

def can_box_move(c, d, grid):
    boxpos = get_box(c, grid)
    assert boxpos is not None

    newpos = boxpos + d
    if is_wall(newpos, grid) or is_wall(newpos + 1j, grid):
        return False

    if d.real == 0:  # przesuwamy się poziomo
        to_check = newpos
        if d.imag == 1:
            to_check += 1j
        return empty(to_check, grid) or can_box_move(to_check, d, grid)
    else:  # przesuwamy się pionowo
        if empty(newpos, grid) and empty(newpos + 1j, grid):
            return True
        for x in [0, 1j]:
            tomovepos = get_box(newpos + x, grid)
            if tomovepos is None:
                continue

            if not can_box_move(tomovepos, d, grid):
                return False
        return True

Funkcja move_boxes jest nieco prostsza bo zakłada już że wszystkie pudełka można przesunąć (design by contract):

def move_boxes(c, d, grid):
    boxpos = get_box(c, grid)
    if boxpos is None:
        return

    newpos = boxpos + d
    if d.real == 0:
        tomovepos = newpos
        if d.imag == 1:
            tomovepos += 1j
        move_boxes(tomovepos, d, grid)
    else:
        for x in [0, 1j]:
            tomovepos = newpos + x
            move_boxes(tomovepos, d, grid)

    grid[newpos] = BOX
    del grid[boxpos]

Testy przechodzą, możemy teraz to wszystko zintegrować

Finalne rozwiązanie

import sys

WALL = 0
BOX = 1


DIR_MAP = {"<": complex(0, -1), "^": complex(-1, 0), ">": 1j, "v": 1}


def parse(data: str):
    grid2d, dirs = data.split("\n\n")
    robot = -1

    grid = dict()
    for i, row in enumerate(grid2d.split("\n")):
        for j, c in enumerate(row):
            comp = complex(i, j * 2)
            if c == "#":
                grid[comp] = WALL
                grid[comp + 1j] = WALL
            elif c == "O":
                grid[comp] = BOX
            elif c == "@":
                robot = comp

    dirs = [DIR_MAP[d] for d in dirs.replace("\n", "")]

    assert robot != -1
    return grid, dirs, robot


def render(grid, robot):
    for i in range(10):
        for j in range(20):
            c = complex(i, j)
            if c == robot:
                print("@", end="")
            elif c not in grid:
                if c - 1j in grid and grid[c - 1j] == BOX:
                    print("]", end="")
                else:
                    print(".", end="")
            elif grid[c] == WALL:
                print("#", end="")
            else:
                print("[", end="")
        print()


def can_box_move(c, d, grid):
    boxpos = get_box(c, grid)
    assert boxpos is not None

    newpos = boxpos + d
    if is_wall(newpos, grid) or is_wall(newpos + 1j, grid):
        return False

    if d.real == 0:
        to_check = newpos
        if d.imag == 1:
            to_check += 1j
        return empty(to_check, grid) or can_box_move(to_check, d, grid)
    else:
        if empty(newpos, grid) and empty(newpos + 1j, grid):
            return True
        for x in [0, 1j]:
            tomovepos = get_box(newpos + x, grid)
            if tomovepos is None:
                continue

            if not can_box_move(tomovepos, d, grid):
                return False
        return True


def empty(c, grid):
    return c not in grid and not (c - 1j in grid and grid[c - 1j] == BOX)


def is_wall(c, grid):
    return c in grid and grid[c] == WALL


def get_box(c, grid):
    if c in grid and grid[c] == BOX:
        return c
    if c - 1j in grid and grid[c - 1j] == BOX:
        return c - 1j
    return None


def move_boxes(c, d, grid):
    boxpos = get_box(c, grid)
    if boxpos is None:
        return

    newpos = boxpos + d
    if d.real == 0:
        tomovepos = newpos
        if d.imag == 1:
            tomovepos += 1j
        move_boxes(tomovepos, d, grid)
    else:
        for x in [0, 1j]:
            tomovepos = newpos + x
            move_boxes(tomovepos, d, grid)

    grid[newpos] = BOX
    del grid[boxpos]


if __name__ == "__main__":
    data = sys.stdin.read()
    grid, dirs, robot = parse(data)

    for it, d in enumerate(dirs):
        new = robot + d
        if empty(new, grid):
            robot = new
            continue

        if is_wall(new, grid):
            continue

        if can_box_move(new, d, grid):
            move_boxes(new, d, grid)
            robot = new

    result = 0
    for c in grid:
        if grid[c] == BOX:
            result += c.real * 100 + c.imag

    print(result)

Podsumowanie

Czuję że rozwiązanie zostało nieco przekomplikowane, możliwe że gdybym zamiast zapisywać obiekty w słowniku zapisywał je normalnie w gridzie 2d to kod byłby trochę prostszy, ale myślę że wydajnościowo udało się uzyskać całkiem dobry wynik.

$ time python3 part2.py < input.txt
python3 part2.py < input.txt  0.02s user 0.01s system 96% cpu 0.031 total

Zadanie#

Rozwiązanie#

Parsowanie#

Przesuwanie#

Liczenie wyniku#

Część druga#

Pomysł#

Implementacja#

Finalne rozwiązanie#

Podsumowanie#