Боря программирует

XOR Linked Tree



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



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



Понятно, что нужно просто запустить bfs/dfs из корня и сохранить расстояние до каждой вершины. Во время bfs нужно уметь обходить всех соседей текущей вершины. Для этого придется исходный список ребер превратить в N списков (для каждой вершины список ее соседей). Обычно N списков хранятся как Vec<Vec<usize>>, а это значит N отдельных аллокаций в heap, каждая из которых в среднем маленького размера, а такие обращения к памяти работают долго. На моем компьютере на каком-то дереве из миллиона вершин такой код работает 154мс.



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



А теперь расскажу алгоритм, который работает 34мс. Для каждой вершины сохраним два числа — количество ее соседей, а также XOR всех ее соседей. Вместо того, чтобы запускать bfs, вначале найдем топологическую сортировку дерева. Будем постепенно удалять листья из дерева, пока не останется только корень. Чтобы узнать, что вершина лист — нужно проверить, что у нее ровно один сосед (и мы знаем какой, он записан в XOR). Когда удаляем лист, нужно обновить XOR и количество соседей у его единственного соседа. После того как все вершины удалены, можно пройтись по ним в обратном порядке и посчитать расстояния до корня.



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

fn get_heights_xor_tree(edges: &[(usize, usize)]) -> Vec<usize> {

    let n = edges.len() + 1;

    let mut cnt_neighbours = vec![0; n];

    let mut xor_neighbours = vec![0; n];

    for &(u, v) in edges {

        cnt_neighbours[u] += 1;

        cnt_neighbours[v] += 1;

        xor_neighbours[u] ^= v;

        xor_neighbours[v] ^= u;

    }

    let mut queue = Vec::with_capacity(n);

    let mut parent = vec![n; n];

    parent[0] = 0;

    for mut v in 0..n {

        while cnt_neighbours[v] == 1 && parent[v] == n {

            queue.push(v);

            let p = xor_neighbours[v];

            parent[v] = p;

            cnt_neighbours[p] -= 1;

            xor_neighbours[p] ^= v;

            v = p;

        }

    }

    assert_eq!(queue.len(), n - 1);

    let mut h = vec![0; n];

    for &v in queue.iter().rev() {

        h[v] = h[parent[v]] + 1;

    }

    h

}