July 28, 2022

Генерируем полимино

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

После конца контеста я вроде бы понял, что не все так страшно, так что давайте обсудим какие есть варианты.

Hardcode

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

В качестве упражнения я потом попытался это сделать. Если вам не лень, то можете проверить свою внимательность и понять, забыл ли я что-то :)

Может лучше кодом?

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

  • перебрать все возможные закраски квадрата 4х4 (их всего 2^16)
  • проверить, что фигура внутри состоит из не более чем 4 клеток
  • проверить, что фигура связная

У этого метода есть проблема, что одну и ту же фигуру мы сгенерируем несколько раз. Например, прямоугольник 1х4 мы сгенерируем 4 раза (как каждую отдельную строку).

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

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

Немного деталей реализации. Будем хранить закрашенные клетки как один инт. Будем считать, что клетка (row, col) закрашена, если в этом инте проставлен (row*4+col)-й бит.

Также будем считать, что у нас уже есть написанная структура данных Dsu, которая умеет поддерживать множества. У dsu есть операция unite, которая принимает два элемента. Если они уже лежат в одном множестве, то возвращает false, а иначе объединяет множества и возвращает true.

Код получается примерно такой:

fn gen_figures_mask(max_cnt: usize) {
    let id = |row: usize, col: usize| row * max_cnt + col;

    let first_row_mask = (1 << max_cnt) - 1;
    let first_col_mask = (0..max_cnt).map(|row| 1 << id(row, 0)).sum::<i32>();

    for mask in 1i32..(1 << (max_cnt * max_cnt)) {
        if mask.count_ones() as usize > max_cnt
            || (mask & first_row_mask) == 0
            || (mask & first_col_mask) == 0
        {
            continue;
        }
        let mut dsu = Dsu::new(max_cnt * max_cnt);
        let mut num_comps = mask.count_ones() as usize;
        for r in 0..max_cnt {
            for c in 0..max_cnt {
                let id1 = id(r, c);
                if (1 << id1) & mask != 0 {
                    if r + 1 < max_cnt && ((1 << id(r + 1, c)) & mask) != 0 {
                        if dsu.unite(id1, id(r + 1, c)) {
                            num_comps -= 1;
                        }
                    }
                    if c + 1 < max_cnt && ((1 << id(r, c + 1)) & mask) != 0 {
                        if dsu.unite(id1, id(r, c + 1)) {
                            num_comps -= 1;
                        }
                    }
                }
            }
        }
        if num_comps == 1 {
            // TODO: return figure
        }
    }
}

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

Breadth-first search

Другой возможный способ — написать bfs по фигурам. Начинаем с фигурки в одну клетку. А потом каждый раз добавляем новую клетку рядом с уже существующей. Такие фигуры будут автоматически связными.

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

Но проблема одинаковых фигур все еще может возникнуть. Например, мы начали с клетки (0, 0), а потом добавили либо клетку (1, 0), либо клетку (-1, 0). В обоих случаях получилась одинаковая фигура. Чтобы такого не происходило можно считать, что клетка (0, 0) должна всегда присутствовать в фигуре, и что она должна быть лексикографически меньше всех остальных.

Предположим, что у нас уже есть структура Point, которая хранит пару (x, y), а так же определена операция + на них.

Тогда код будет примерно таким:

#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Figure(Vec<Point>);

impl Figure {
    pub fn new(mut pts: Vec<Point>) -> Self {
        pts.sort();
        Self(pts)
    }
}

const SHIFTS_4: [Point; 4] = [
    Point { x: 0, y: 1 },
    Point { x: 0, y: -1 },
    Point { x: 1, y: 0 },
    Point { x: -1, y: 0 },
];

pub fn gen_figures(max_cnt: usize) -> Vec<Figure> {
    let start = Figure::new(vec![Point::ZERO]);
    let mut seen = HashSet::new();
    let mut queue = VecDeque::new();
    seen.insert(start.clone());
    queue.push_back(start);
    while let Some(fig) = queue.pop_back() {
        if fig.0.len() == max_cnt {
            continue;
        }
        for p in fig.0.iter() {
            for shift in SHIFTS_4.iter() {
                let new_point = p + shift;
                if new_point > Point::ZERO && !fig.0.contains(&new_point) {
                    let next_figure = Figure::new([fig.0.clone(), vec![new_point]].concat());
                    if !seen.contains(&next_figure) {
                        seen.insert(next_figure.clone());
                        queue.push_back(next_figure);
                    }
                }
            }
        }
    }
    seen.into_iter().collect()
}

В таком коде уже явно меньше какой-то битовой магии и сложнее ошибиться.

Bonus

Первый способ работал достаточно быстро, потому что 2^16 это не так много. Но если бы нужно было генерировать фигурки с большим количеством клеток, он бы работал значительно дольше.

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

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

Если вы дочитали до сюда, то подписывайтесь на мой канал https://t.me/bminaiev_blog