Kilor 31 мар в 10:40

SQL HowTo: работаем с массивами (Advent of Code 2024, Day 23: LAN Party)

Простой

7 мин

1.2K

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

Кейс

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

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

Применяем простые операции над массивами, чтобы определить связность графов.

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

Advent of Code 2024, Day 23: LAN Party

--- День 23: LAN-вечеринка ---

Пока Историки бродят по безопасной зоне в штаб-квартире Easter Bunny, вы натыкаетесь на плакаты о запланированной на сегодня LAN-вечеринке! Возможно, вам удастся ее найти; вы подключаетесь к ближайшему порту передачи данных и загружаете карту локальной сети (ваши входные данные для головоломки).

Карта сети предоставляет список всех соединений между двумя компьютерами. Например:

kh-tc
qp-kh
de-cg
ka-co
yn-aq
qp-ub
cg-tb
vc-aq
tb-ka
wh-tc
yn-cg
kh-ub
ta-co
de-co
tc-td
tb-wq
wh-td
ta-ka
td-qp
aq-cg
wq-ub
ub-vc
de-ta
wq-aq
wq-vc
wh-yn
ka-de
kh-ta
co-tc
wh-qp
tb-vc
td-yn

Каждая строка текста на карте сети представляет собой одно соединение; строка kh-tc представляет собой соединение между компьютером с именем kh и компьютером с именем tc. Соединения не являются направленными; tc-kh означало бы то же самое.

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

В этом примере имеются 12 таких наборов из трех соединенных между собой компьютеров:

aq,cg,yn
aq,vc,wq
co,de,ka
co,de,ta
co,ka,ta
de,ka,ta
kh,qp,ub
qp,td,wh
tb,vc,wq
tc,td,wh
td,wh,yn
ub,vc,wq

Если Главный Историк здесь, и он на LAN-вечеринке, лучше всего узнать об этом сразу. Вы почти уверены, что имя его компьютера начинается с t, поэтому рассматривайте только наборы из трех компьютеров, где имя хотя бы одного компьютера начинается с t. Это сужает список до 7 наборов из трех взаимосвязанных компьютеров:

co,de,ta
co,ka,ta
de,ka,ta
qp,td,wh
tb,vc,wq
tc,td,wh
td,wh,yn

Найдите все наборы из трех соединенных между собой компьютеров. Сколько из них содержат хотя бы один компьютер с именем, начинающимся с t?

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

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

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

В приведенном выше примере самый большой набор компьютеров, которые все подключены друг к другу, состоит из co, de, ka, и ta. Каждый компьютер в этом наборе имеет подключение к каждому другому компьютеру в наборе:

ka-co
ta-co
de-co
ta-ka
de-ta
ka-de

На плакатах LAN-вечеринки говорится, что пароль для входа на LAN-вечеринку - это имена каждого компьютера на LAN-вечеринке, отсортированные в алфавитном порядке, а затем соединенные запятыми. (Люди, организующие LAN-вечеринку, явно являются кучкой зануд.) В этом примере пароль будет co,de,ka,ta.

Какой пароль для входа на LAN-вечеринку?

Часть 1

Как обычно, начнем с разбора исходных данных в набор пар связанных компьютеров:

WITH src AS (
  SELECT
    regexp_matches(
      $$
kh-tc
qp-kh
de-cg
ka-co
yn-aq
qp-ub
cg-tb
vc-aq
tb-ka
wh-tc
yn-cg
kh-ub
ta-co
de-co
tc-td
tb-wq
wh-td
ta-ka
td-qp
aq-cg
wq-ub
ub-vc
de-ta
wq-aq
wq-vc
wh-yn
ka-de
kh-ta
co-tc
wh-qp
tb-vc
td-yn
$$
    , '([a-z]{2})-([a-z]{2})'
    , 'g'
    ) pair
)

  pair
 text[]
{kh,tc}
{qp,kh}
{de,cg}
{ka,co}
{yn,aq}
{qp,ub}
...

Чтобы найти 3 взаимно соединенных компьютера воспользуемся двумя соединениями, запретив в каждом соединение с "меньшей" парой:

SELECT
  *
FROM
  src p1
JOIN
  src p2
    ON p2.pair > p1.pair AND
    p2.pair && p1.pair              -- пары имеют общие элементы
JOIN
  src p3
    ON p3.pair > p2.pair AND
    p3.pair <@ (p1.pair || p2.pair) -- оба элемента третьей пары присутствуют в первых двух

{ka,co} | {ta,co} | {ta,ka}
{vc,aq} | {wq,aq} | {wq,vc}
{kh,ub} | {qp,kh} | {qp,ub}
{de,co} | {ka,co} | {ka,de}
{de,co} | {de,ta} | {ta,co}
{tc,td} | {wh,tc} | {wh,td}
{td,qp} | {wh,qp} | {wh,td}
{aq,cg} | {yn,aq} | {yn,cg}
{ub,vc} | {wq,ub} | {wq,vc}
{de,ta} | {ka,de} | {ta,ka}
{tb,vc} | {tb,wq} | {wq,vc}
{td,yn} | {wh,td} | {wh,yn}

Добавим вычисление условия "какое-то из имен начинается на t" и просуммируем количество записей, которые будут ему соответствовать:

SELECT
  sum(cond::integer)                -- подсчитываем удовлетворяющие условию
FROM
  src p1
JOIN
  src p2
    ON p2.pair > p1.pair AND
    p2.pair && p1.pair              -- пары имеют общие элементы
JOIN
  src p3
    ON p3.pair > p2.pair AND
    p3.pair <@ (p1.pair || p2.pair) -- оба элемента третьей пары присутствуют в первых двух
, LATERAL (
    SELECT
      bool_or(name ^@ 't') cond     -- какое-то из имен среди всех пар начинается на t
    FROM
      unnest(p1.pair || p2.pair || p3.pair) name
  ) T;

Соединим в итоговый запрос:

WITH src AS (
  SELECT
    regexp_matches(
      $$
kh-tc
qp-kh
de-cg
ka-co
yn-aq
qp-ub
cg-tb
vc-aq
tb-ka
wh-tc
yn-cg
kh-ub
ta-co
de-co
tc-td
tb-wq
wh-td
ta-ka
td-qp
aq-cg
wq-ub
ub-vc
de-ta
wq-aq
wq-vc
wh-yn
ka-de
kh-ta
co-tc
wh-qp
tb-vc
td-yn
$$
    , '([a-z]{2})-([a-z]{2})'
    , 'g'
    ) pair
)
SELECT
  sum(cond::integer)                -- подсчитываем удовлетворяющие условию
FROM
  src p1
JOIN
  src p2
    ON p2.pair > p1.pair AND
    p2.pair && p1.pair              -- пары имеют общие элементы
JOIN
  src p3
    ON p3.pair > p2.pair AND
    p3.pair <@ (p1.pair || p2.pair) -- оба элемента третьей пары присутствуют в первых двух
, LATERAL (
    SELECT
      bool_or(name ^@ 't') cond     -- какое-то из имен среди всех пар начинается на t
    FROM
      unnest(p1.pair || p2.pair || p3.pair) name
  ) T;

Поскольку мы не применяли вообще никаких упрощений и оптимизаций (типа предварительного формирования только пар содержащих имена "на t"), решая задачу "в лоб" самым простым и очевидным способом, то полный "кубический" перебор (два соединения) 3080 пар и отфильтровкой почти 150M неподходящих под условия записей занимает 72 секунды:

Часть 2

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

Сначала соберем данные, с кем вообще связан каждый "напрямую" (то есть является первым в указанной или "обратной" паре):

, paired AS (
  SELECT
    pair[1] src                     -- кто связан
  , array_agg(DISTINCT pair[2]) dst -- с кем связан (уникализированный набор)
  FROM
    (
      TABLE src                 -- все пары в прямом порядке
    UNION ALL
      SELECT
        ARRAY[pair[2], pair[1]] -- все пары в обратном порядке
      FROM
        src
    )
  GROUP BY
    1
)

src  | dst
text | text[]
aq   | {cg,vc,wq,yn}
cg   | {aq,de,tb,yn}
co   | {de,ka,ta,tc}
de   | {cg,co,ka,ta}
ka   | {co,de,ta,tb}
kh   | {qp,ta,tc,ub}
qp   | {kh,td,ub,wh}
ta   | {co,de,ka,kh}
tb   | {cg,ka,vc,wq}
tc   | {co,kh,td,wh}
td   | {qp,tc,wh,yn}
ub   | {kh,qp,vc,wq}
vc   | {aq,tb,ub,wq}
wh   | {qp,tc,td,yn}
wq   | {aq,tb,ub,vc}
yn   | {aq,cg,td,wh}

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

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

, r AS (
  SELECT
    ARRAY[src] party -- начинаем со всех одинарных групп
  FROM
    paired
UNION ALL
  SELECT
    party || src
  FROM
    r
  JOIN
    paired
      ON src > party[array_length(party, 1)] AND -- новый элемент больше последнего
      party <@ dst                               -- все уже собранные присутствуют среди соединенных с новым
)

party
{aq}
{cg}
{co}
{de}
{ka}
{kh}
{qp}
{ta}
{tb}
{tc}
{td}
{ub}
{vc}
{wh}
{wq}
{yn}
{aq,cg}
{cg,de}
{co,de}
{co,ka}
{de,ka}
{kh,qp}
{co,ta}
{de,ta}
{ka,ta}
{kh,ta}
{cg,tb}
{ka,tb}
{co,tc}
{kh,tc}
{qp,td}
{tc,td}
{kh,ub}
{qp,ub}
{aq,vc}
{tb,vc}
{ub,vc}
{qp,wh}
{tc,wh}
{td,wh}
{aq,wq}
{tb,wq}
{ub,wq}
{vc,wq}
{aq,yn}
{cg,yn}
{td,yn}
{wh,yn}
{co,de,ka}
{co,de,ta}
{co,ka,ta}
{de,ka,ta}
{kh,qp,ub}
{qp,td,wh}
{tc,td,wh}
{aq,vc,wq}
{tb,vc,wq}
{ub,vc,wq}
{aq,cg,yn}
{td,wh,yn}
{co,de,ka,ta}

Остается лишь преобразовать самую первую из самых длинных групп в строку:

SELECT
  array_to_string(party, ',')
FROM
  r
ORDER BY
  array_length(party, 1) DESC -- берем самую длинную группу
, party                       -- сортируем по алфавиту, если их вдруг окажется несколько
LIMIT 1;

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

WITH RECURSIVE src AS (
  SELECT
    regexp_matches(
      $$
kh-tc
qp-kh
de-cg
ka-co
yn-aq
qp-ub
cg-tb
vc-aq
tb-ka
wh-tc
yn-cg
kh-ub
ta-co
de-co
tc-td
tb-wq
wh-td
ta-ka
td-qp
aq-cg
wq-ub
ub-vc
de-ta
wq-aq
wq-vc
wh-yn
ka-de
kh-ta
co-tc
wh-qp
tb-vc
td-yn
$$
    , '([a-z]{2})-([a-z]{2})'
    , 'g'
    ) pair
)
, paired AS (
  SELECT
    pair[1] src                     -- кто связан
  , array_agg(DISTINCT pair[2]) dst -- с кем связан (уникализированный набор)
  FROM
    (
      TABLE src                 -- все пары в прямом порядке
    UNION ALL
      SELECT
        ARRAY[pair[2], pair[1]] -- все пары в обратном порядке
      FROM
        src
    )
  GROUP BY
    1
)
, r AS (
  SELECT
    ARRAY[src] party -- начинаем со всех одинарных групп
  FROM
    paired
UNION ALL
  SELECT
    party || src
  FROM
    r
  JOIN
    paired
      ON src > party[array_length(party, 1)] AND -- новый элемент больше последнего
      party <@ dst                               -- все уже собранные присутствуют среди соединенных с новым
)
SELECT
  array_to_string(party, ',') -- превращаем отсортированный по построению массив в строку
FROM
  r
ORDER BY
  array_length(party, 1) DESC -- берем самую длинную группу
, party                       -- сортируем по алфавиту, если их вдруг окажется несколько
LIMIT 1;

Решая гораздо более сложную задачу (ища группы не только из 3, но из максимального количества соединенных компьютеров), мы вышли практически на то же время - 76 секунд:

UPD: Альтернативный вариант из комментариев - без рекурсии и всего за ~180мс (спасибо @z_sergза идею с подсчетом количества связей внутри группы - для полного графа количество ребер должно быть равно треугольному числу n(n - 1)/2).

Часть 1 (снова)

Раз сложную задачу можно решить за то же время, то простую должно быть можно существенно быстрее!

Достаточно ограничить размер собираемых наборов "каждый с каждым" 3 элементами - это и будут наши искомые "триплеты":

-- ...
, r AS (
  -- ...
  WHERE
    array_length(party, 1) < 3 -- ограничиваем длину набора 3 элементами
)
SELECT
  sum(cond::integer)
FROM
  r
, LATERAL (
    SELECT
      bool_or(unnest ^@ 't') cond
    FROM
      unnest(party)
  ) T
WHERE
  array_length(party, 1) = 3; -- проверяем только наборы длины 3

В таком варианте решение становится в 50 раз быстрее и занимает всего лишь чуть больше 1.5 секунд:

UPD: Альтернативный вариант из комментариев - 60мс.

Теги:

Хабы: