Задача 1

Дано корневое дерево и число $$$k$$$. Для каждой вершины $$$u$$$ найти кол-во вершин в поддереве вершины $$$u$$$ отдаленных от $$$u$$$ на расстояние не более $$$k$$$ рёбер.

Задача 2

Дано дерево и число $$$k$$$. Для каждой вершины $$$u$$$ найти кол-во вершин во всем дереве отдаленных от $$$u$$$ на расстояние не более $$$k$$$ рёбер. (источник) (похожая задача)

В этом блоге рассмотрим решение Задачи 1 за $$$O(N)$$$. При этом используется забавная структура — RQ на stack-е с запросами $$$O(1)$$$ (по сути это будет префикс-сумма).

Вообще Задача 1 возникла при попытке (видимо не удачной) решить Задачу 2, для которой есть решения за $$$O(N \ln N)$$$, но я пока не знаком с этими темами и не хочу напутать:

• через Centroid Decomposition

• для похожей задачи через SplayTree и прочее (см. комментарии про задачу D) $$$ $$$

Решение

Решать Задачу 1 будем через DFS. При обходе будем хранить стек чисел, каждое значение которого будет соответствовать ответу для вершины определенного поддерева. Длинна стека $$$h$$$ будет соответствовать глубине рекурсии в DFS.

• При входе в новую вершину, увеличиваем значения в интервале стека $$$[h-k, h)$$$ на единицу. И push-ем $$$1$$$ в стек.

• При выходе из вершины, сохраняем последнее значения стека в вершину соответствующего поддерева. И затем делам pop в стеке.

Оказывается эти операции с стеком можно выполнять за $$$O(1)$$$:

• Для увеличения на интервале выполняем: diff[l] -= delta; diff[r] += delta;

• При pop() выполняем "пропихивание": diff[l-1] += diff[l];

struct RQ_stack {
    int l = 0;
    vector<int> data;
    vector<int> diff;  // нумерация с 1

    RQ_stack() { diff.resize(1); }  // фиктивное значение в начале
    size_t size() { return l; }

    void push(int x = 0) {
        if( l == data.size() ) {
            data.push_back({});
            diff.push_back({});
        }
        data[l]   = x;
        diff[l+1] = 0;
        l++;
    }

    int pop() {
        l--;
        diff[l] += diff[l+1];
        return data[l] + diff[l+1];
    }

    void increase(int l, int r, int delta = 1) {
        diff[l] -= delta;
        diff[r] += delta;
    }
};

vector<vector<int>> edges;
vector<int> subtree_results_k;

int dfs_k;
RQ_stack yoko;
void dfs(int i, int prev, vector<int>& subtree_results) {
    int l = yoko.size();    
    yoko.increase( max(0, l-dfs_k), l );
    yoko.push(1);

    for(auto j : edges[i]) {
        if( j == prev ) continue;
        
        dfs(j, i, subtree_results);
    }
    subtree_results[i] = yoko.pop();
}

После поисков нашёл задачу где используется что-то схожее — хранение пути от корня в DFS: F. Два поддерева (решение).

Хотелось бы воспользоваться rerootin-ом и с помощью Задачи 1 решить Задачу 2. Однако тогда для корня надо знать результаты для дистанций { $$$k, k-1$$$ }, чтобы подсчитать результат для дистанции { $$$k$$$ } у ребёнка. В итоге у ребёнка теряется результат для дистанции { $$$k-1$$$ }, и как его можно оптимально восстанавливать не понятно.

Удивительно, но мне как-то мало попалось статьей об этих задачах, надеюсь этот блог будет полезен!