Fun with Postgres Looped Functions and Linear Progressions

Disclaimer

This article will contain spoilers both on how I solved 2022 Day 21's challenge "Monkey Math" 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.

AOC Day 21

Tech used:

• The file_fdw extension to read the input
• Functions such as regexp_substr
• Unlogged tables

As always, we will use file_fdw to put our text input into a virtual Postgres table:

CREATE EXTENSION if not exists file_fdw;

CREATE SERVER if not exists aoc2022 foreign data wrapper file_fdw;

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

CREATE FOREIGN TABLE aoc_day21 (id text, action text)
  SERVER aoc2022 options(filename '/tmp/aoc2022.day21.input',
--  SERVER aoc2022 options(filename '/tmp/aoc2022.day21.testinput',
  FORMAT 'csv', DELIMITER ':');

AOC Day 21 - Part One

The puzzle directions are odd but parseable:

Each monkey is given a job: either to yell a specific number or to yell
the result of a math operation. All of the number-yelling monkeys
know their number from the start; however, the math operation monkeys
need to wait for two other monkeys to yell a number, and those two
other monkeys might also be waiting on other monkeys.

We don't speak monkey, but the elephants we freed in the previous rounds do. This puzzle is pretty straightforward. First, let's pull apart the text strings in the puzzle, which look like this:

cgrb: gzwb * rcfd
gfbz: bwgp - qlfm
jrbf: 2
gvvg: rjch + tjdp
vwsh: grwp * ddsv
tpwb: 1

We will separate the data in each line and store one monkey per row in a new unlogged table. As each row is guaranteed to have a colon, we declared the foreign table as a csv with a delimiter of a colon, which saves us a step. But we still need to break apart the other items into specific columns. Some simple regular expression functions can help us do this:

CREATE UNLOGGED TABLE puzzle (
  id     TEXT,
  number BIGINT,
  lefty  TEXT,
  action TEXT,
  righty TEXT
);
/* Fill sparsely, as we will be updating this table a lot */
ALTER TABLE puzzle SET (autovacuum_enabled = off, fillfactor = 20);

INSERT INTO puzzle SELECT id
  ,CASE WHEN action ~ '\d'
    THEN regexp_substr(action, '(\d+)')::BIGINT ELSE null END
  ,CASE WHEN action !~ '\d'
    THEN regexp_substr(action, '\w+') ELSE null END
  ,CASE WHEN action !~ '\d'
    THEN regexp_substr(action, '[+*/-]') ELSE null END
  ,CASE WHEN action !~ '\d'
    THEN ltrim(regexp_substr(ltrim(action), ' (\w+)')) ELSE null END
FROM aoc_day21;

For each line, we examine if it has a number in it or not. If it does, we need to extract the monkey name ("id") and the number it yells out. If there is no number, we need to extract what other monkeys are involved, and what the mathematical symbol is. Afterwards, the table looks like this (we used \pset NULL ☃ in our .psqlrc file to produce a better null indicator)

  id  | number | lefty | action | righty
------+--------+-------+--------+--------
 cgrb |      ☃ | gzwb  | *      | rcfd
 gfbz |      ☃ | bwgp  | -      | qlfm
 jrbf |      2 | ☃     | ☃      | ☃
 gvvg |      ☃ | rjch  | +      | tjdp
 vwsh |      ☃ | grwp  | *      | ddsv
 tpwb |      1 | ☃     | ☃      | ☃

A function will be used to walk through monkey by monkey, apply any math that is needed, and keep running through until finally the monkey named "root" says a number, which will be our solution.

CREATE FUNCTION riddle_me_this()
  RETURNS BIGINT language plpgsql AS $$
DECLARE
  myrec RECORD; first BIGINT; second BIGINT;
BEGIN

LOOP
  /* Walk through and solve every monkey that has a left value. Order does not matter */
  FOR myrec IN SELECT * FROM puzzle WHERE number IS NULL LOOP
      /* Record the number yelled by the first monkey we are listening to */
      SELECT INTO first  p.number FROM puzzle p WHERE id = myrec.lefty;
      IF first IS NULL THEN continue; END IF;
      /* Record the number yelled by the second monkey */
      SELECT INTO second  p.number FROM puzzle p WHERE id = myrec.righty;
      IF second IS NULL THEN continue; END IF;

      /* At this point, we have numbers from two other monkeys, so perform an action */
      UPDATE puzzle SET number =
        CASE WHEN myrec.action = '-' THEN first - second
             WHEN myrec.action = '+' THEN first + second
             WHEN myrec.action = '*' THEN first * second
             WHEN myrec.action = '/' THEN first / second
        END
      WHERE id = myrec.id;

      IF myrec.id = 'root' THEN
        RETURN number FROM puzzle WHERE id = myrec.id;
      END IF;
    END LOOP;
  END LOOP;
END
$$;

We are just about ready to run the function. As we are doing a lot of lookups based on the "id" column, we want to create an index for it:

CREATE INDEX monkey_id ON puzzle(id);

Finally, we analyze the table, turn on timing, and run the function to get the correct answer:

ANALYZE puzzle; /* This helps a lot! */
\timing on
SELECT riddle_me_this();  /* Took 630ms on my system for a 1619 line input file */

riddle_me_this
-----------------
 158731561459602

AOC Day 21 - Part Two

For the second part of the puzzle, we need to figure out what number to feed into the "humn" row such that the "root" row will eventually have the same left and right values. To achieve this, we'll loop through a few times. Based on the previous day's puzzles, this will require a LOT of rounds, so we'll start with a guess of one trillion and then compute how far off we are. So first, we run with a guess of one, then of one trillion, and then compute the differences. All these monkeys are forming a simple linear progression, so we can quickly narrow in at that point until we get matching "root" numbers and have our answer.

CREATE FUNCTION i_am_humn()
  RETURNS BIGINT language plpgsql AS $$
DECLARE
  round INT = 0; myrec RECORD;
  first BIGINT; oldfirst BIGINT; second BIGINT;
  human_value BIGINT; changerate FLOAT;
  trillion BIGINT = 1_000_000_000;
BEGIN
<<outer>> LOOP
  round = round + 1;
  IF round = 1 THEN
    human_value = 1;
  ELSIF round = 2 THEN
    human_value = trillion;
  END IF;

  RAISE INFO '-> Round %: starting with human_value of %', round,
    to_char(human_value, 'FM999G999G999G999G999');

  /* reset to initial state */
  TRUNCATE TABLE puzzle;
  INSERT INTO puzzle SELECT id
    ,CASE WHEN action ~ '\d'  THEN
      regexp_substr(action, '(\d+)')::BIGINT ELSE null END
    ,CASE WHEN action !~ '\d' THEN
      regexp_substr(action, '\w+') ELSE null END
    ,CASE WHEN action !~ '\d' THEN
      regexp_substr(action, '[+*/-]') ELSE null END
    ,CASE WHEN action !~ '\d' THEN
      ltrim(regexp_substr(ltrim(action), ' (\w+)')) ELSE null END
  FROM aoc_day21;

  <<inner>> LOOP

    FOR myrec IN SELECT * FROM puzzle WHERE number IS  NULL LOOP
      SELECT INTO first   p.number FROM puzzle p WHERE id = myrec.lefty;
      IF first IS NULL THEN continue; END IF;
      SELECT INTO second  p.number FROM puzzle p WHERE id = myrec.righty;
      IF second IS NULL THEN continue; END IF;

      /* Discard the original values for "humn" as the goal is for us to provide them */
      IF myrec.lefty  = 'humn' THEN first = human_value; END IF;
      IF myrec.righty = 'humn' THEN second = human_value; END IF;

      UPDATE puzzle SET number =
        CASE WHEN myrec.action = '-' THEN first - second
             WHEN myrec.action = '+' THEN first + second
             WHEN myrec.action = '*' THEN first * second
             WHEN myrec.action = '/' THEN first / second
        END
      WHERE id = myrec.id;

      /* If this is monkey "root" AND the values are the same, we have finished */
      IF myrec.id = 'root' THEN
        IF first = second THEN RETURN human_value; END IF;
        EXIT inner;
      END IF;
    END LOOP;
  END LOOP;

  /* If this is our second run, see how far a trillion numbers has pushed us */
  IF z = round THEN
    changerate = (first-oldfirst) / trillion::float;
  END IF;

  /* Once we know how fast we change based on the input, we can refine our guess */
  IF round >= 2 THEN
    IF first-second < 0 THEN
      human_value = floor(human_value - abs((first-second)/changerate));
    ELSE
      human_value = human_value + abs((first-second)/changerate);
    END IF;
  END IF;

  oldfirst = first;

END LOOP;
END
$$;

There are two things in the function above to make things easier for us humans. First, the use of tochar(human_value, 'FM999G999G999G999G999') rather than just human_value ensures that a bigint like 3769668748355 gets output as 3,769,668,748,355. Second, one of my favorite features of Postgres 16 is the ability to write long numbers in a friendly manner. That's why instead of the confusing BIGINT = 1000000000 we can now simply write BIGINT = 1_000_000_000, Those of you copy and pasting this into an earlier-than-16 version will see: ERROR: trailing junk after numeric literal at or near "1*". Let's run the function:

SET client_min_messages = 'INFO';
SELECT i_am_humn();

INFO:  -> Round 1: starting with human_val of 1
INFO:  -> Round 2: starting with human_val of 1,000,000,000
INFO:  -> Round 3: starting with human_val of 3,769,668,748,355
INFO:  -> Round 4: starting with human_val of 3,769,668,716,709
   i_am_humn
---------------
 3769668716709

This produced the answer in 2.7 seconds, which I am going to call a win. Hopefully this is the last we see of the monkeys this year!