Задача 1

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

Рассмотрим решение за $$$O(N)$$$. При этом возникнет забавная структура — RQ на stack-е с запросами $$$O(1)$$$ (по сути префикс-сумма). $$$ $$$

Решение

Запускаем DFS. При обходе будем хранить стек чисел: вершине пути на глубине $$$h$$$ соответствует $$$h$$$-ый индекс стека.

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

• При выходе из вершины, сохраняем последнее значения стека в вершину соответствующего поддерева (это если что подсчитан ответ: кол-во вершин на расстояние не более $$$k$$$). И затем делам 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();
}

Можно слегка модифицировать (аккумулировать сколько к текущему моменту вершин ниже этой высоты) и решать: Дано корневое дерево. Надо оффлайн отвечать на запросы: $$$(u,k)$$$ — количество вершин поддерева $$$u$$$, отдаленных от $$$u$$$ на расстоянии не более чем $$$k$$$ рёбер ($$$k$$$ для каждого запроса своё).

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

Если рассматривать не только вершины в поддереве, а все вершины в дереве, то такую задачу можно решить за $$$O(N \ln N)$$$:

Задача 2

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

• через Centroid Decomposition

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

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

соре за бред, писалось 15 месяцев назад, для увеличения шума в интернете оставлю этот блог (+ постарался уменьшить ерунды)