SQL HowTo: оконные функции (Advent of Code 2024, Day 22: Monkey Market)

В этой челлендж-серии статей попробуем использовать PostgreSQL как среду для решения задач Advent of Code 2024.

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

Используем оконные функции, чтобы вычислить "третью производную".

• Решение Day 1: Historian Hysteria (регулярные выражения и условная агрегация)

• Решение Day 2: Red-Nosed Reports (логические агрегаты)

• Решение Day 3: Mull It Over ("чистые" регулярки)

• Решение Day 4: Ceres Search (работа с массивами)

• Решение Day 5: Print Queue (поиск в словаре и массивах, сортировка "пузырьком")

• Решение Day 6: Guard Gallivant (рекурсивные циклы и их контроль)

• Решение Day 7: Bridge Repair ("экспоненциальная" рекурсия)

• Решение Day 8: Resonant Collinearity (генерация и подсчет уникальных комбинаций)

• Решение Day 9: Disk Fragmenter (оптимизируем рекурсию)

• Решение Day 10: Hoof It (поиск "в ширину" внутри цикла)

• Решение Day 11: Plutonian Pebbles (агрегация внутри рекурсии)

• Решение Day 12: Garden Groups (волновой алгоритм и подсчет границ)

• Решение Day 13: Claw Contraption (пошагово решаем СЛУ)

• Решение Day 14: Restroom Redoubt (находим "елочку" с помощью центра масс)

• Решение Day 15: Warehouse Woes (играем в сокобан с помощью json-карты и типа point)

• Решение Day 16: Reindeer Maze (укрощаем рекурсию в лабиринте)

• Решение Day 17: Chronospatial Computer (подбираем значение ветвлением)

• Решение Day 18: RAM Run (поиск пути и дихотомия)

• Решение Day 19: Linen Layout (динамическое программирование)

• Решение Day 20: Race Condition (кратчайший путь "туда и обратно" и его самосоединение)

• Решение Day 21: Keypad Conundrum (моделирование против подсчета)

• Решение Day 22: Monkey Market (оконные функции)

• Решение Day 23: LAN Party (работаем с массивами)

• Решение Day 24: Crossed Wires (немного двоичной логики)

• Решение Day 25: Code Chronicle (простой финал с агрегатами)

---

Оригинальная постановка задачи и ее перевод:

Advent of Code 2024, Day 22: Monkey Market

--- День 22: Обезьяний рынок ---

Пока вы все телепортировались вглубь джунглей, обезьяна крадет устройство Историков! Вам нужно будет вернуть его, пока Историки ищут Шефа.

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

Вы не уверены, как работает Monkey Exchange Market, но один из Историков чувствует проблему и приходит на помощь. По-видимому, они уже некоторое время изучают этих обезьян и расшифровали их секреты.

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

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

Часть о секретах буквальна, объясняет Историк. Каждый покупатель создает псевдослучайную последовательность секретных чисел, где каждый секрет выводится из предыдущего.

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

• Вычислите результат умножения секретного числа на 64. Затем смешайте этот результат с секретным числом. Наконец, обрежьте секретное число.

• Рассчитайте результат деления секретного числа на 32. Округлите результат до ближайшего целого числа. Затем смешайте этот результат с секретным числом. Наконец, обрежьте секретное число.

• Вычислите результат умножения секретного числа на 2048. Затем смешайте этот результат с секретным числом. Наконец, обрежьте секретное число.

Каждый этап вышеописанного процесса включает смешивание и обрезку:

• Чтобы смешать значение с секретным числом, вычислите побитовое XOR заданного значения и секретного числа. Затем секретное число становится результатом этой операции. (Если секретное число равно 42 и вы смешали его с секретным числом 15, секретное число станет 37.)

• Чтобы обрезать секретное число, вычислите значение секретного числа по модулю 16777216. Затем секретное число станет результатом этой операции. (Если секретное число равно 100000000, и вы собираетесь обрезать его, секретное число станет 16113920.)

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

Итак, если у покупателя есть секретное число 123, то следующие десять секретных чисел этого покупателя будут такими:

15887950
16495136
527345
704524
1553684
12683156
11100544
12249484
7753432
5908254

Каждый покупатель использует свой собственный секретный номер при выборе цены, поэтому важно уметь предсказывать последовательность секретных номеров для каждого покупателя. К счастью, исследование Историка раскрыло начальный секретный номер каждого покупателя (ваш пазл). Например:

1
10
100
2024

Этот список описывает начальный секретный номер четырех разных покупателей секретных убежищ на Monkey Exchange Market. Если вы можете смоделировать секретные номера от каждого покупателя, вы сможете предсказать все их будущие цены.

За один день у каждого покупателя есть время сгенерировать 2000 новых секретных номеров. В этом примере для каждого покупателя их начальный секретный номер и 2000-й секретный номер, который они сгенерируют, следующие:

1: 8685429
10: 4700978
100: 15273692
2024: 8667524

Сложение 2000-х секретных номеров всех покупателей дает 37327623.

Для каждого покупателя смоделируйте создание 2000 новых секретных чисел. Какова сумма 2000-х секретных чисел, сгенерированных каждым покупателем?

--- Часть вторая ---

Конечно, секретные числа - это не цены, которые предлагает каждый покупатель! Это было бы смешно. Вместо этого цены, которые предлагает покупатель, - это просто последние цифры каждого из их секретных чисел.

Таким образом, если покупатель начинает с секретного числа 123, то первые десять цен этого покупателя будут следующими:

3 (от 123)
0 (от 15887950)
6 (от 16495136)
5 (и т.д.)
4
4
6
4
4
2

Эта цена — количество бананов, которое покупатель предлагает в обмен на вашу информацию о новом укрытии. Однако вы все еще не говорите на языке обезьян, поэтому вы не можете вести переговоры с покупателями напрямую. Историк немного говорит, но недостаточно, чтобы вести переговоры; вместо этого он может попросить другую обезьяну вести переговоры от вашего имени.

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

Таким образом, если покупатель начинает с секретного числа 123, то первые десять секретных чисел этого покупателя, цены и связанные с ними изменения будут следующими:

     123: 3 
15887950: 0 (-3)
16495136: 6 (6)
  527345: 5 (-1)
  704524: 4 (-1)
 1553684: 4 (0)
12683156: 6 (2)
11100544: 4 (-2)
12249484: 4 (0)
 7753432: 2 (-2)

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

В этом коротком примере, в пределах только этих первых нескольких цен, самая высокая цена будет 6, поэтому было бы неплохо дать обезьяне инструкции, которые заставят ее продать в это время. Первая 6 происходит только после двух изменений, поэтому нет способа дать обезьяне инструкцию продать в этот момент, но второй раз 6 появляется после изменений -1,-1,0,2. Таким образом, если вы дадите обезьяне эту последовательность изменений, она будет ждать, пока впервые не увидит эту последовательность, а затем немедленно продаст вашу информацию о месте укрытия по текущей цене, выиграв вам 6 бананов.

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

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

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

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

1
2
3
2024

Существует много последовательностей из четырех изменений цен, которые вы могли бы рассказать обезьяне, но для этих четырех покупателей последовательность, которая принесет вам больше всего бананов, это -2,1,-1,3. Используя эту последовательность, обезьяна совершит следующие продажи:

• Для покупателя с начальным секретным числом 1 изменения -2,1,-1,3 в первую очередь встречаются, когда цена составляет 7.

• Для покупателя с первоначальным секретом 2 изменения -2,1,-1,3 в первую очередь встречаются, когда цена составляет 7.

• Для покупателя с начальным секретом 3 последовательность изменений -2,1,-1,3 не встречается в первых 2000 изменениях.

• Для покупателя, начинающего с 2024, изменения -2,1,-1,3 в первую очередь встречаются, когда цена составляет 9.

Итак, попросив обезьяну продать в первый раз, когда цены каждого покупателя пойдут вниз 2, затем вверх 1, затем вниз 1, затем вверх 3, вы получите 23 (7 + 7 + 9) банана!

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

Часть 1

В первой части нас просят вычислить сумму 2000-х сгенерированных из исходных по определенным правилам чисел. Традиционно, получим эти самые исходные "секретные" числа регулярными выражениями:

WITH RECURSIVE src AS (
  SELECT
    line[1]::integer secret
  FROM
    regexp_matches(
      $$
1
10
100
2024
$$
    , '[^\r\n]+(?=$|[\r\n])'
    , 'g'
    ) line
)

Затем реализуем три указанные последовательные операции вычисления следующего секретного числа:

, r AS (
  SELECT
    0 i
  , secret src
  , secret
  FROM
    src
UNION ALL
  SELECT
    i + 1
  , src
  , secret3
  FROM
    r
  , LATERAL (
      SELECT
        r.secret secret0
      , 0xFFFFFF wrap -- "по модулю" 16777216 (2^24)
    ) T0
  , LATERAL (
      SELECT
        (secret0 << 6) # secret0 & wrap secret1 -- x64, смешать и обрезать
    ) T1
  , LATERAL (
      SELECT
        (secret1 >> 5) # secret1 & wrap secret2 -- :32, смешать и обрезать
    ) T2
  , LATERAL (
      SELECT
        (secret2 << 11) # secret2 & wrap secret3 -- x2048, смешать и обрезать
    ) T3
  WHERE
    i < 2000
)

В принципе, можно придумать и более эффективный вариант вычисления, поскольку все битовые операции происходят только в нижних 24 битах (за счет "обрезки"), но нам хватит и такой производительности.

Осталось лишь найти сумму последних, 2000-х, значений:

SELECT
  sum(secret)
FROM
  (
    TABLE r
    ORDER BY
      i DESC          -- берем с последнего шага рекурсии ...
    FETCH NEXT 1 ROWS
    WITH TIES         -- ... все записи с тем же значением сортировки
  ) T;

Собственно, вот и все решение первой части:

WITH RECURSIVE src AS (
  SELECT
    line[1]::integer secret
  FROM
    regexp_matches(
      $$
1
10
100
2024
$$
    , '[^\r\n]+(?=$|[\r\n])'
    , 'g'
    ) line
)
, r AS (
  SELECT
    0 i
  , secret
  FROM
    src
UNION ALL
  SELECT
    i + 1
  , secret3
  FROM
    r
  , LATERAL (
      SELECT
        r.secret secret0
      , 0xFFFFFF wrap -- "по модулю" 16777216 (2^24)
    ) T0
  , LATERAL (
      SELECT
        (secret0 << 6) # secret0 & wrap secret1 -- x64, смешать и обрезать
    ) T1
  , LATERAL (
      SELECT
        (secret1 >> 5) # secret1 & wrap secret2 -- :32, смешать и обрезать
    ) T2
  , LATERAL (
      SELECT
        (secret2 << 11) # secret2 & wrap secret3 -- x2048, смешать и обрезать
    ) T3
  WHERE
    i < 2000
)
SELECT
  sum(secret)
FROM
  (
    TABLE r
    ORDER BY
      i DESC          -- берем с последнего шага рекурсии ...
    FETCH NEXT 1 ROWS
    WITH TIES         -- ... все записи с тем же значением сортировки
  ) T;

Вычисления для 2K+ реальных секретных чисел занимают меньше 4 секунд:

Часть 2

Во второй части нас просят сначала вычислить секретные числа, затем взять от каждого последнюю цифру ("первая производная"), затем вычислить изменение между соседними цифрами ("вторая производная"), а потом найти цепочку из 4 последовательных изменений, дающую максимальную сумму значений этих цифр-цен.

Вот тут как раз сложность условия очень легко превращается в решение на SQL с помощью оконных функций.

Сначала немного модифицируем решение из первой части, чтобы сохранить ID каждого покупателя:

WITH RECURSIVE src AS (
  SELECT
    id -- ID покупателя
  , line[1]::integer secret
  FROM
    regexp_matches(
      $$
1
2
3
2024
$$
    , '[^\r\n]+(?=$|[\r\n])'
    , 'g'
    )
      WITH ORDINALITY T(line, id) -- используем автонумерацию строк
)
, r AS (
  SELECT
    id
  , 0 i
  , secret
  FROM
    src
UNION ALL
  SELECT
    id -- сохраняем исходный ID покупателя на всех шагах
  , i + 1
  , secret3
  FROM
    r
  , LATERAL (
      SELECT
        r.secret secret0
      , 0xFFFFFF wrap
    ) T0
  , LATERAL (
      SELECT
        (secret0 << 6) # secret0 & wrap secret1
    ) T1
  , LATERAL (
      SELECT
        (secret1 >> 5) # secret1 & wrap secret2
    ) T2
  , LATERAL (
      SELECT
        (secret2 << 11) # secret2 & wrap secret3
    ) T3
  WHERE
    i < 2000
)

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

• вычислим "цену" (последнюю цифру) от каждого рассчитанного нами секретного числа
  secret % 10

• найдем изменение цены от предыдущего по порядку значения для того же клиента
  price - lag(price) OVER(PARTITION BY id ORDER BY i)

• получим последовательности для каждых 4 следующих друг за другом строк
  array_agg(diff) OVER(PARTITION BY id ORDER BY i ROWS BETWEEN 3 PRECEDING AND CURRENT ROW)

• уберем "неработающие" последовательности для первых трех цен (у них в первой позиции стоит NULL)
  seq[1] IS NOT NULL

• для всех остальных последовательностей оставим по каждому клиенту только первый по порядку раз, когда такая последовательность встретилась
  SELECT DISTINCT ON(seq, id) ... ORDER BY seq, id, i

• наконец, сгруппировав в разрезе последовательностей, найдем максимум суммы цен
  sum(price) ... GROUP BY seq ORDER BY sum DESC LIMIT 1

Или на SQL:

SELECT
  sum(price)
FROM
  (
    SELECT DISTINCT ON(seq, id) -- оставляем только первое появление последовательности для клиента
      *
    FROM
      (
        SELECT
          *
        , array_agg(diff) OVER(PARTITION BY id ORDER BY i ROWS BETWEEN 3 PRECEDING AND CURRENT ROW) seq -- собираем последовательность из 4 изменений
        FROM
          (
            SELECT
              *
            , price - lag(price) OVER(PARTITION BY id ORDER BY i) diff -- вычисляем изменение цены
            FROM
              (
                SELECT
                  *
                , secret % 10 price -- вычисляем цену
                FROM
                  r
              ) T
          ) T
      ) T
    WHERE
      seq[1] IS NOT NULL -- исключаем "неработающие" последовательности
    ORDER BY
      seq, id, i
  ) T
GROUP BY
  seq
ORDER BY
  sum DESC -- находим максимум суммы цен
LIMIT 1;

Соберем все вместе:

WITH RECURSIVE src AS (
  SELECT
    id -- ID покупателя
  , line[1]::integer secret
  FROM
    regexp_matches(
      $$
1
2
3
2024
$$
    , '[^\r\n]+(?=$|[\r\n])'
    , 'g'
    )
      WITH ORDINALITY T(line, id) -- используем автонумерацию строк
)
, r AS (
  SELECT
    id
  , 0 i
  , secret
  FROM
    src
UNION ALL
  SELECT
    id -- сохраняем исходный ID покупателя на всех шагах
  , i + 1
  , secret3
  FROM
    r
  , LATERAL (
      SELECT
        r.secret secret0
      , 0xFFFFFF wrap
    ) T0
  , LATERAL (
      SELECT
        (secret0 << 6) # secret0 & wrap secret1
    ) T1
  , LATERAL (
      SELECT
        (secret1 >> 5) # secret1 & wrap secret2
    ) T2
  , LATERAL (
      SELECT
        (secret2 << 11) # secret2 & wrap secret3
    ) T3
  WHERE
    i < 2000
)
SELECT
  sum(price)
FROM
  (
    SELECT DISTINCT ON(seq, id) -- оставляем только первое появление последовательности для клиента
      *
    FROM
      (
        SELECT
          *
        , array_agg(diff) OVER(PARTITION BY id ORDER BY i ROWS BETWEEN 3 PRECEDING AND CURRENT ROW) seq -- собираем последовательность из 4 изменений
        FROM
          (
            SELECT
              *
            , price - lag(price) OVER(PARTITION BY id ORDER BY i) diff -- вычисляем изменение цены
            FROM
              (
                SELECT
                  *
                , secret % 10 price -- вычисляем цену
                FROM
                  r
              ) T
          ) T
      ) T
    WHERE
      seq[1] IS NOT NULL -- исключаем "неработающие" последовательности
    ORDER BY
      seq, id, i
  ) T
GROUP BY
  seq
ORDER BY
  sum DESC -- находим максимум суммы цен
LIMIT 1;

В этот раз для получения результата придется подождать (все-таки у нас больше 2000 покупателей, да по 2000 цен у каждого), но тоже не слишком долго - около 36 секунд: