Kilor 5 часов назад

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

Простой

6 мин

451

Блог компании ТензорPostgreSQL*SQL*Занимательные задачкиАлгоритмы*

Кейс

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

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

Advent of Code 2024, Day 11: Plutonian Pebbles

--- День 11: Плутонианская галька ---

Древняя цивилизация на Плутоне была известна своей способностью манипулировать пространством-временем, и пока Историки исследуют ее бесконечные коридоры, вы заметили странный набор камней, бросающих вызов законам физики.

На первый взгляд они кажутся обычными камнями: они расположены в идеально прямую линию, и на каждом камне выгравирован номер.

Самое странное, что каждый раз, когда вы моргаете, камни меняются.

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

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

Если на камне выгравирован номер 0, его заменяют камнем с выгравированным номером 1.
Если на камне выгравировано число, имеющее четное количество цифр, его заменяют двумя камнями. Левая половина цифр гравируется на новом левом камне, а правая половина цифр гравируется на новом правом камне. (Новые числа не содержат дополнительных ведущих нулей: 1000станут камнями 10и 0.)
Если ни одно из других правил не применимо, камень заменяется новым; на новом камне гравируется номер старого камня, умноженный на 2024.

Как бы ни менялись камни, их порядок сохраняется, и они остаются на своей идеально прямой линии.

Как будут меняться камни, если вы будете продолжать моргать, глядя на них? Вы записываете число, выгравированное на каждом камне в линии (ваш вход для головоломки).

Если у вас есть композиция из пяти камней с выгравированными числами 0 1 10 99 999и вы моргнете один раз, камни изменятся следующим образом:

Первый камень, 0, становится камнем с маркировкой 1.
Второй камень, 1, умножается на 2024 и становится 2024.
Третий камень, 10, разделится на пару камней, обозначенных как 1 и 0.
Четвертый камень, 99, разделится на два, обозначенных номером 9.
Пятый камень, 999, заменен камнем с маркировкой 2021976.

Таким образом, после того, как вы моргнете один раз, ваши пять камней превратятся в композицию из семи камней с выгравированными числами 1 2024 1 0 9 9 2021976.

Вот более длинный пример:

Исходное состояние:
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

В этом примере, после 6 морганий у вас будет 22 камня. После 25 морганий у вас будет 55312 камней!

Рассмотрите расположение камней перед собой. Сколько камней у вас будет после того, как вы моргнете 25 раз?

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

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

Для этого реализуем достаточно простой рекурсивный запрос, используя для хранения номера на камне тип 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.

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

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

...
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;

Посмотрим, что у нас теперь получается в плане выполнения запроса при данных из второй (сложной) части задачи - исходный набор 965842 9159 3372473 311 0 6 86213 48 и 75 морганий. Воспользуемся нашим сервисом визуализации explain.tensor.ru:

Полное время выполнения заняло чуть меньше 400мс, а ответ - 218279375708592.

Теги:

Хабы:

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

--- День 11: Плутонианская галька ---

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

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

Публикации

Информация