SQL HowTo: агрегация внутри рекурсии (Advent of Code 2024, Day 11: Plutonian Pebbles)

Сегодня посмотрим на примере задачки из Advent of Code зачем и как можно обойти ошибку aggregate functions are not allowed in a recursive query's recursive term, возникающую при попытке агрегировать какие-то данные внутри шага рекурсии на PostgreSQL — «если нельзя, но очень хочется, то можно».

На очередную задачу, для которой получается красивое SQL-решение, не далее чем вчера навел @Olegas — приведу перевод ее условия:

Исходное состояние: 125 17  После 1 моргания: 253000 1 7  После 2 морганий: 253 0 2024 14168  После 3 морганий: 512072 1 20 24 28676032  После 4 морганий: 512 72 2024 2 0 2 4 2867 6032  После 5 морганий: 1036288 7 2 20 24 4048 1 4048 8096 28 67 60 32  После 6 морганий: 2097446912 14168 4048 2 0 2 4 40 48 2024 40 48 80 96 2 8 6 7 6 0 3 2

Наивный вариант

Для начала реализуем вариант «как слышится, так и пишется» и проверим, что получающийся у нас набор камней совпадает с приведенным в примере.

Для этого реализуем достаточно простой рекурсивный запрос, используя для хранения номера на камне тип numeric, поскольку не знаем заранее, насколько длинным он может оказаться:

WITH RECURSIVE src AS (   SELECT     '125 17' nums -- исходный набор камней   , 6 blinks      -- количество морганий ) , r AS (   SELECT     0 blink                                       -- количество прошедших морганий   , regexp_split_to_table(nums, ' ')::numeric num -- развернули строку в набор чисел   FROM     src UNION ALL   SELECT     blink + 1 -- увеличиваем счетчик морганий   , unnest(   -- "разворачиваем" элементы массива       CASE         -- правило #1         WHEN num = 0 THEN       ARRAY[1]         -- правило #2         WHEN length(num::text) % 2 = 0 THEN           ARRAY[             left(num::text, length(num::text) >> 1)           , right(num::text, length(num::text) >> 1)           ]::numeric[]         -- правило #3         ELSE           ARRAY[num * 2024]       END     )   FROM     r   , src   WHERE     blink < blinks -- ограничение по количеству морганий ) SELECT   num FROM   r , src WHERE   blink = blinks; -- отбираем только камни на последнем шаге

Из интересных приемов тут стоит отметить «размножение» записей в зависимости от условия внутри рекурсии с помощью unnest - CASE - ARRAY.

Поскольку мы все сделали правильно, то должны получить следующий результат:

 num  numeric 2097446912      14168       4048          2          0          2          4         40         48       2024         40         48         80         96          2          8          6          7          6          0          3          2

Понятно, что для подсчета итогового количества достаточно немного изменить конец запроса:

... SELECT   count(*) FROM   r , src WHERE   blink = blinks;

Подсчет количества

Однако, в пределе, количество камней может расти по экспоненте, увеличиваясь вдвое каждый (ну, или почти каждый) раз, приводя нас к задаче о зернах на шахматной доске.

Понятно, что при ограничении на 25 морганий, в пределе мы можем получить лишь 2 ^ 25 = 33 554 432 камней, формируя по ходу рекурсии вдвое больше записей. Для современных ресурсов это не слишком большие значения — но что если количество морганий нам увеличат до 100?..

Для более эффективного решения заметим, что по условию задачи о самом итоговом расположении камней нас не спрашивают — лишь об их количестве. При этом одинаково маркированные камни на каждом следующем шаге будут вести себя одинаково, многократно повторяясь в списке — например #0, #2, #40, #48 в примере выше.

Поэтому нам достаточно учитывать лишь количество уникально пронумерованных камней, суммируя количество по каждому из номеров на каждом шаге алгоритма, а в итоге просуммировать их общее количество.

Попробуем выразить это на SQL:

WITH RECURSIVE src AS (   SELECT     '125 17' nums   , 6 blinks ) , r AS (   SELECT     0 blink   , regexp_split_to_table(nums, ' ')::numeric num   , 1::numeric qty -- количество камней с такой маркировкой, исходно по 1   FROM     src UNION ALL   SELECT     blink + 1   , unnest(       CASE         WHEN num = 0 THEN       ARRAY[1]         WHEN length(num::text) % 2 = 0 THEN           ARRAY[             left(num::text, length(num::text) >> 1)           , right(num::text, length(num::text) >> 1)           ]::numeric[]         ELSE           ARRAY[num * 2024]       END     )   , sum(qty) -- суммарное количество   FROM     r   , src   WHERE     blink < blinks   GROUP BY     1, 2 -- группируем по номеру шага и камня ) SELECT   num , qty FROM   r , src WHERE   blink = blinks;

Увы, нас сразу ждет разочарование в виде ошибки:

ERROR:  aggregate functions are not allowed in a recursive query's recursive term LINE 26:   , sum(qty)              ^ ********** Error **********  ERROR: aggregate functions are not allowed in a recursive query's recursive term SQL state: 42P19 Character: 513

Но решение есть!

Давайте реализуем шаг рекурсии внутри вспомогательной CTE, а агрегировать будем уже по ней:

WITH RECURSIVE src AS (   SELECT     '125 17' nums   , 6 blinks ) , r AS (   SELECT     0 blink   , regexp_split_to_table(nums, ' ')::numeric num   , 1::numeric qty   FROM     src UNION ALL   ( -- это нужная скобка!     WITH tmp AS (       SELECT         blink + 1 blink       , unnest(           CASE             WHEN num = 0 THEN               ARRAY[1]             WHEN length(num::text) % 2 = 0 THEN               ARRAY[                 left(num::text, length(num::text) >> 1)               , right(num::text, length(num::text) >> 1)               ]::numeric[]             ELSE               ARRAY[num * 2024]           END         ) num       , qty       FROM         r       , src       WHERE         blink < blinks     )     SELECT       blink     , num     , sum(qty) qty     FROM       tmp -- агрегируем по "нерекурсивной" CTE     GROUP BY       1, 2   ) ) SELECT   num -- выводим каждый номер , qty -- ... и количество таких камней FROM   r , src WHERE   blink = blinks;

Замечу, что дополнительная скобка вокруг рекурсивной части запроса теперь обязательна из-за использования WITH, иначе возникнет ошибка:

ERROR:  syntax error at or near "WITH" LINE 15:     WITH tmp AS (

Теперь список камней из примера выглядит так после 6 итераций:

  num      | qty   numeric  | numeric         40 |       2          4 |       1         96 |       1          2 |       4       4048 |       1      14168 |       1       2024 |       1 2097446912 |       1         48 |       2          3 |       1          0 |       2          7 |       1          6 |       2          8 |       1         80 |       1

Осталось лишь просуммировать общее количество — и решение готово:

WITH RECURSIVE src AS (   SELECT     '125 17' nums   , 6 blinks ) , r AS (   SELECT     0 blink   , regexp_split_to_table(nums, ' ')::numeric num   , 1::numeric qty   FROM     src UNION ALL   (     WITH tmp AS (       SELECT         blink + 1 blink       , unnest(           CASE             WHEN num = 0 THEN               ARRAY[1]             WHEN length(num::text) % 2 = 0 THEN               ARRAY[                 left(num::text, length(num::text) >> 1)               , right(num::text, length(num::text) >> 1)               ]::numeric[]             ELSE               ARRAY[num * 2024]           END         ) num       , qty       FROM         r       , src       WHERE         blink < blinks     )     SELECT       blink     , num     , sum(qty) qty     FROM       tmp     GROUP BY       1, 2   ) ) SELECT   sum(qty) count FROM   r , src WHERE   blink = blinks;