Advent of Code - Day 9

AOC Day 9

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

The plpgsql 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 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()
  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;
   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;
    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 colcolumns (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
-------
   152

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:

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()
  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)
             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.


SELECT drawsnake();

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;

Loading terminal...