How to Solve Advent of Code 2022 Using Postgres - Day 9

Spoiler Alert!

This article will contain spoilers both on how I solved 2022 Day 9's challenge "Rope Bridge" using SQL, as well as general ideas on how to approach the problem. I recommend trying to solve it yourself first, using your favorite language. This article is delayed from the actual puzzle's release. Also note that my solutions are seldom going to be the "best" solutions - they were solved as quickly as possible, and these articles will show my first solutions, with some minor reformatting and cleaning up.

Hands on Tutorial

We've also loaded a tutorial for Day 9's challenge if you want to try it with a pre-loaded data set.

AOC Day 9

This challenge was a fun one. Some of the items used:

• file_fdw, a built-in foreign data wrapper for Postgres
• The PL/pgSQL language
• ASCII art - with colors!
• Upsert and merge

The data file is very simple, and consists of a direction and an amount. We need to track the head of a "rope" and the tail of the "rope". And by rope, we really mean snake, because we are going to ASCII animate a snake!

R 4
U 4
L 3
D 1
R 4
D 1
L 5
R 2

Our usual setup, including a custom schema to stick things in. We will add a delimiter call when creating our foreign table as it is already nicely formatted for us:

CREATE EXTENSION if not exists file_fdw;

CREATE SERVER if not exists aoc2022 foreign data wrapper file_fdw;

DROP SCHEMA if exists snake CASCADE;
CREATE SCHEMA snake;
SET search_path = snake;

DROP FOREIGN TABLE if exists aoc_day9;

CREATE FOREIGN TABLE aoc_day9 (dir char, far smallint)
  SERVER aoc2022 options(filename '/tmp/aoc2022.day9.input', delimiter ' ');

DROP SEQUENCE if exists aoc;
CREATE SEQUENCE if not exists aoc;

We want to put everything into a table, in which each location on a grid is represented by a single database row. We also add in a collection of indexes for later use, via the DO command to add some simple looping help:

DROP TABLE if exists snakegrid;
CREATE UNLOGGED TABLE snakegrid (
  row int, col int,
  head bool not null default false,
  middle text not null default '',
  tailcount int not null default 0
);
CREATE UNIQUE INDEX snakey ON snakegrid(row,col);
CREATE INDEX snakehead on snakegrid(head);

DO $$
BEGIN FOR x IN 1..9 LOOP
  EXECUTE format($i$ create index snakemiddle%s on snakegrid(row,col) where middle ~ '%s' $i$, x, x);
END LOOP;
END
$$;

Next we populate the table:

INSERT INTO snakegrid SELECT x,y from generate_series(1,20) x, generate_series(1,20) y;
UPDATE snakegrid SET head=true, middle='123456789', tailcount=1 WHERE row=15 and col=12;

This is a 20x20 grid. The snake may wander outside the grid, but we have an UPSERT coming up to handle that situation. For fun, let's make a function to draw the grid:

CREATE or replace FUNCTION drawsnake() --version 1
  returns void language plpgsql AS $$
DECLARE
 myrec record; r int; myrow text; output text = '';
 hrow int; hcol int; trow int; tcol int; trails int;
BEGIN
  PERFORM pg_sleep(0.1);
  FOR r IN SELECT DISTINCT row FROM snakegrid ORDER BY 1 LOOP
    myrow = '';
    FOR myrec IN SELECT * FROM snakegrid WHERE row=r ORDER BY col LOOP
      myrow = format('%s %s', myrow,
        CASE WHEN myrec.head THEN E'\033[31mH\033[0m'
             WHEN myrec.middle ~ '9' THEN E'\033[33mT\033[0m'
             WHEN myrec.tailcount>0 THEN E'\033[32m*\033[0m' ELSE ' ' END);
    END LOOP;
    output = output || myrow || chr(10);
 END LOOP;
 SELECT INTO hrow,hcol row,col FROM snakegrid WHERE head;
 SELECT INTO trow,tcol row,col FROM snakegrid WHERE middle ~ '9';
 SELECT INTO trails COUNT(*) FROM snakegrid WHERE tailcount > 0;
 RAISE NOTICE '% % % % HEAD=%,%  TAIL=%,%  TRAILS=%', E'\033c', chr(10), output, chr(10),
   E'\033[31m'|| format('%2s', hrow), format('%2s', hcol) || E'\033[0m',
   E'\033[33m'|| format('%2s', trow), format('%2s', tcol) || E'\033[0m',
   E'\033[32m'||trails||E'\033[0m';

END
$$;

There is a lot going on in this little function. Basically, we walk through each cell in the grid, and output which item is in it. If the head of the snake is there, we output the letter H. We are also going to start adding some ANSI color codes, just to make things a little more pretty. If the tail appears, as indicated by the number 9 in the "middle" column, we output a yellow T. If the tail has already been in that area (which is what the challenge asks us to check), then we output a green #. At the end, we do our usual screen-clearing trick, and add some tracking at the bottom to show current position of the snake head, the current position of the snake tail, and the total number of places the tail has visited.

Now all we need is a way to solve the puzzle. Another function to the rescue. This one takes three parameters - a direction (R,L,U,D), the number of spaces to move, and a boolean indicating if we want to redraw the screen each time (e.g. animate things).

CREATE or replace FUNCTION movesnake(dir char, moves int, drawsnake bool default false)
  returns void language plpgsql as $$
DECLARE
  hrow int; hcol int; trow int; tcol int; newtrow int; newtcol int;
BEGIN
  SELECT INTO hrow,hcol row,col FROM snakegrid WHERE head LIMIT 1;
  FOR i IN 1..moves LOOP
    UPDATE snakegrid SET head=false WHERE row=hrow AND col=hcol;
    IF 'R' = dir THEN hcol = hcol + 1; END IF;
    IF 'L' = dir THEN hcol = hcol - 1; END IF;
    IF 'U' = dir THEN hrow = hrow - 1; END IF;
    IF 'D' = dir THEN hrow = hrow + 1; END IF;
    INSERT INTO snakegrid VALUES (hrow,hcol,true)
      ON CONFLICT (row,col) DO UPDATE SET head=true;

   -- Move the tail if needed
   SELECT INTO trow,tcol row,col FROM snakegrid WHERE middle ~ '9' LIMIT 1;
   newtrow = trow; newtcol = tcol;

   IF    hrow = trow THEN newtcol = newtcol + (hcol - tcol)/2;
   ELSIF hcol = tcol THEN newtrow = newtrow + (hrow - trow)/2;
   ELSE
     IF ABS(hrow-trow)>=2 OR ABS(hcol-tcol)>=2
       THEN newtcol = tcol + (CASE WHEN hcol>tcol THEN 1 ELSE -1 END);
            newtrow = trow + (CASE WHEN hrow>trow THEN 1 ELSE -1 END);
     END IF;
   END IF;

    -- Only update if we need to
    IF trow <> newtrow OR tcol <> newtcol THEN
      UPDATE snakegrid SET middle = '' WHERE row=trow AND col=tcol;
      INSERT INTO snakegrid(row,col,middle,tailcount) VALUES (newtrow,newtcol,'9',1)
        ON CONFLICT (row,col) DO UPDATE SET middle = '9', tailcount = snakegrid.tailcount + 1;
    END IF;

    IF drawsnake THEN PERFORM drawsnake(); END IF;
  END LOOP;
  RETURN;
END
$$;

In the function above, we first locate the single row that contains the current location of the snake's head, and grab information about it. Then we loop for however many moves we are performing this invocation. In each loop, we remove the head marker from its current location, then stick it into one of the nearby grid locations. We do this with an upsert("update or insert"), as we don't know if the grid we are moving to already exists or not. First the INSERT is attempted, but if it gets denied due to a conflict on the row and col columns (recall the unique index), we instead UPDATE the existing row:

INSERT INTO snakegrid VALUES (hrow,hcol,true)
  ON CONFLICT (row,col) DO UPDATE SET head=true;

Next we move on to (and move) the tail of the snake. There are specific rules given in the puzzle as to how it follows the head, which we follow by changing the newtcol and newtrow as needed. Once that is done, we move the tail in another upsert similar to the head, but we also increment the tailcount variable (or set it to 1 if this is a new row). Finally, we optionally call the drawsnake() function.

Once the two functions are in place, we can invoke them directly from the foreign table: no need for any CTEs on this one!

SELECT movesnake(dir,far,false) from aoc_day9;

SELECT count(*) FROM snakegrid WHERE tailcount > 0;

 count
-------
   13

For part two, the snake grows a lot larger! Rather than just a head and tail, it now has 8 segments between the two tails (and thus our tail becomes segment 9). Each of these segments follows the previous one, so the head moves, then segment 1 moves based on the head, then segment 2 moves based on segment 1, etc. A new function based on the previous one should do the trick:

CREATE or replace FUNCTION movebigsnake(dir char, moves int, drawsnake bool default false)
  returns void language plpgsql as $$
DECLARE
  hrow int; hcol int; trow int; tcol int; newtrow int; newtcol int;
  lastrow smallint; lastcol smallint;
  myrec record;
BEGIN
  SELECT INTO hrow,hcol row,col FROM snakegrid WHERE head LIMIT 1;
  FOR i IN 1..moves LOOP
    UPDATE snakegrid SET head=false WHERE row=hrow AND col=hcol;
    IF 'R' = dir THEN hcol = hcol + 1; END IF;
    IF 'L' = dir THEN hcol = hcol - 1; END IF;
    IF 'U' = dir THEN hrow = hrow - 1; END IF;
    IF 'D' = dir THEN hrow = hrow + 1; END IF;
    INSERT INTO snakegrid VALUES (hrow,hcol,true)
      ON CONFLICT (row,col) DO UPDATE SET head=true;

    -- Move each segment in turn, ending with the tail
    lastrow = hrow; lastcol = hcol;
    FOR segment IN 1..9 LOOP
      SELECT INTO trow,tcol row,col FROM snakegrid WHERE middle ~ segment::text  LIMIT 1;
      newtrow = trow; newtcol = tcol;

      IF    lastrow = trow THEN newtcol = newtcol + (lastcol - tcol)/2;
      ELSIF lastcol = tcol THEN newtrow = newtrow + (lastrow - trow)/2;
      ELSE
        IF ABS(lastrow-trow)>=2 OR ABS(lastcol-tcol)>=2
           THEN newtcol = tcol + (CASE WHEN lastcol>tcol THEN 1 ELSE -1 END);
                newtrow = trow + (CASE WHEN lastrow>trow THEN 1 ELSE -1 END);
        END IF;
      END IF;

      -- Only update if we need to
      IF trow <> newtrow OR tcol <> newtcol THEN
        UPDATE snakegrid SET middle = replace(middle, segment::text, '') WHERE row=trow AND col=tcol;

        -- Postgres 14 and older use this:
        -- INSERT INTO snakegrid(row,col,middle,tailcount)
        --   VALUES (newtrow,newtcol,segment,(CASE WHEN segment=9 THEN 1 ELSE 0 END))
        --   ON CONFLICT (row,col) DO UPDATE SET middle = snakegrid.middle || segment,
        --     tailcount = snakegrid.tailcount + (CASE WHEN segment=9 THEN 1 ELSE 0 END);

        MERGE INTO snakegrid s
        USING (SELECT newtrow AS row, newtcol AS col) AS v
        ON s.row=v.row AND s.col=v.col
        WHEN MATCHED THEN
          UPDATE SET middle = middle || segment, tailcount = tailcount + (CASE WHEN segment=9 THEN 1 ELSE 0 END)
        WHEN NOT MATCHED THEN
          INSERT (row,col,middle,tailcount) VALUES (newtrow,newtcol,segment,(CASE WHEN segment=9 THEN 1 ELSE 0 END));
      END IF;

      -- Help out the next loop:
      lastrow = newtrow;
      lastcol = newtcol;

    END LOOP;
    IF drawsnake THEN PERFORM drawsnake(); END IF;

  END LOOP;
  RETURN;
END
$$;

It is mostly the same as the other, but for the second upsert I decided to use the new merge feature of Postgres, which does everything an upsert does and more:

MERGE INTO snakegrid s
  USING (SELECT newtrow AS row, newtcol AS col) AS v ON s.row=v.row AND s.col=v.col
  WHEN MATCHED THEN
    UPDATE SET middle = middle || segment, tailcount = tailcount + (CASE WHEN segment=9 THEN 1 ELSE 0 END)
  WHEN NOT MATCHED THEN
     INSERT (row,col,middle,tailcount) VALUES (newtrow,newtcol,segment,(CASE WHEN segment=9 THEN 1 ELSE 0 END));

For the merge, we tell it how to join the new data with the existing table, and what to do if the row already exists (UPDATE!) and what to do if the rows do not (INSERT!). We use the middle column to maintain a simple text field showing what segments are in that particular spot, as many segments are allowed to co-exists. Each number may appear more than once in the middle column, but that's okay. Finally, we make sure to update our tailcount, but only when the segment is 9 (aka the tail)

Two more tweaks for the drawsnake() function: add in a clause to output the segment numbers directly, and change our viewing "window" so we can track the snake even it moves away from our initial 20 x 20 area:

CREATE or replace FUNCTION drawsnake() -- version 2
  returns void language plpgsql AS $$
DECLARE
 myrec record; r int; myrow text; output text = '';
 hrow int; hcol int; trow int; tcol int; trails int;
BEGIN
  PERFORM pg_sleep(0.1);
  FOR r IN (select greatest(1,row-15) from snakegrid where head) ..
           (select row+15 from snakegrid where head) LOOP
    myrow = '';
    FOR myrec IN SELECT * FROM snakegrid WHERE row=r ORDER BY col LOOP
      myrow = format('%s %s', myrow,
        CASE WHEN myrec.head THEN E'\033[31mH\033[0m'
             WHEN myrec.middle ~ '9' THEN E'\033[33mT\033[0m'
             WHEN length(myrec.middle)>0 THEN right(myrec.middle,1) -- right = most recent addition
             WHEN myrec.tailcount>0 THEN E'\033[32m*\033[0m' ELSE ' ' END);
    END LOOP;
    output = output || myrow || chr(10);
 END LOOP;
 SELECT INTO hrow,hcol row,col FROM snakegrid WHERE head;
 SELECT INTO trow,tcol row,col FROM snakegrid WHERE middle ~ '9';
 SELECT INTO trails COUNT(*) FROM snakegrid WHERE tailcount > 0;
 RAISE NOTICE '% % % % HEAD=%,%  TAIL=%,%  TRAILS=%', E'\033c', chr(10), output, chr(10),
   E'\033[31m'|| format('%2s', hrow), format('%2s', hcol) || E'\033[0m',
   E'\033[33m'|| format('%2s', trow), format('%2s', tcol) || E'\033[0m',
   E'\033[32m'||trails||E'\033[0m';

END
$$;

Finally, we can run our new version and get the answer. Here's a picture of it drawing the test input given on the challenge page:

-- Reset from above:
TRUNCATE TABLE snakegrid;
INSERT INTO snakegrid SELECT x,y from generate_series(1,20) x, generate_series(1,20) y;
UPDATE snakegrid SET head=true, middle='123456789', tailcount=1 WHERE row=15 and col=12;

SELECT movebigsnake(dir,far,true) FROM aoc_day9;

SELECT count(*) FROM snakegrid WHERE tailcount > 0;