Блог пользователя spike1236

Метод виртуальных деревьев
Автор spike1236, история, 7 часов назад, По-русски

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

Введение

Во многих задачах, связанных с деревьями, требуется работать с подмножеством вершин, сохраняя при этом структуру дерева, индуцированную их попарными наименьшими общими предками (LCA). Понятие виртуального дерева позволяет перейти от исходного дерева $$$T$$$ с $$$N$$$ вершинами к подструктуре, размер которой линейно зависит от размера выбранного множества $$$X$$$. Это существенно ускоряет алгоритмы, в частности, для подсчета ДП.

Определение виртуального дерева

Пусть $$$T$$$ — заданное корневое дерево, а $$$X$$$ — некоторое подмножество его вершин. Виртуальное дерево $$$A(X)$$$ определяется следующим образом:

$$$ A(X) = \{ \operatorname{lca}(x,y) \mid x,y \in X \}, $$$

где $$$\operatorname{lca}(x,y)$$$ означает наименьшего общего предка вершин $$$x$$$ и $$$y$$$ в дереве $$$T$$$. В этом дереве между каждой парой вершин проводится ребро, если одна из них является предком другой в $$$T$$$.

Такое определение гарантирует, что в структуру попадают все вершины, «важные» для анализа $$$X$$$ (то есть все LCA для пар вершин из $$$X$$$), а связность наследуется от исходного дерева.

Ключевые свойства и их доказательства

Свойство 1: Оценка по размеру

Утверждение.
Для любого множества $$$X$$$ вершин дерева $$$T$$$ справедливо неравенство:

$$$ \vert A(X) \vert \le 2 \vert X \vert - 1. $$$

Доказательство.
Выберем порядок обхода в глубину (DFS) и рассмотрим вершины множества $$$X$$$, расположенные в порядке обхода:

$$$ x_1, x_2, \dots, x_m. $$$

Обозначим

$$$ l = \operatorname{lca}(x_1, x_2). $$$

Далее необходимо доказать следующую лемму.

Лемма 1. Для любого целого числа $$$n \ge 3$$$, если $$$\operatorname{lca}(x_1,x_n) \neq l$$$, то

$$$ \operatorname{lca}(x_1,x_n) = \operatorname{lca}(x_2,x_n). $$$

Доказательство (от противного).
Предположим, что для некоторого $$$n \ge 3$$$ выполняется $$$\operatorname{lca}(x_1,x_n) \neq l$$$ и одновременно $$$\operatorname{lca}(x_1,x_n) \neq \operatorname{lca}(x_2,x_n)$$$. Пусть $$$k = \operatorname{lca}(x_1,x_n)$$$.
Если $$$k$$$ является предком $$$l$$$, то по определению получаем, что и $$$\operatorname{lca}(x_1,x_n)$$$, и $$$\operatorname{lca}(x_2,x_n)$$$ равны $$$k$$$, что противоречит предположению. Следовательно, $$$k$$$ должно быть потомком $$$l$$$. Но тогда в порядке DFS должен получиться порядок $$$x_1, x_n, x_2$$$, что противоречит выбранному порядку обхода.
Противоречие завершает доказательство леммы.

Возвращаясь к доказательству свойства, заметим, что согласно лемме, каждый LCA, возникающий при последовательном рассмотрении вершин $$$X$$$, равен либо: 1. самой вершине $$$x_1$$$, 2. $$$\operatorname{lca}(x_1,x_2)$$$, либо 3. $$$\operatorname{lca}(x_2,x_n)$$$ для некоторого $$$n$$$.

Таким образом, можно рекурсивно записать:

$$$ A(\{x_1,\dots,x_m\}) = A(\{x_2,\dots,x_m\}) \cup \{ x_1, \operatorname{lca}(x_1,x_2) \}. $$$

Из этого следует, что

$$$ \vert A(X) \vert \le \vert A(\{x_2,\dots,x_m\}) \vert + 2. $$$

Применяя индукцию по размеру множества $$$X$$$, получаем итоговую оценку:

$$$ \vert A(X) \vert \le 2 \vert X \vert - 1. $$$

При этом, если дерево $$$T$$$ является совершенным двоичным деревом, а $$$X$$$ состоит из его листьев, неравенство достигается в точности.

Свойство 2: Представление через последовательные LCA

Утверждение.
Пусть вершины $$$X$$$ упорядочены по порядку обхода DFS: $$$x_1, x_2, \dots, x_m$$$. Тогда

$$$ A(X) = X \cup \{ \operatorname{lca}(x_i, x_{i+1}) \mid 1 \le i < m \}. $$$

Доказательство.
Начнём с рекурсивного представления:

$$$ A(\{x_1,\dots,x_m\}) = A(\{x_2,\dots,x_m\}) \cup \{ x_1, \operatorname{lca}(x_1,x_2) \}. $$$

Раскроем его рекурсивно:

$$$ \begin{aligned} A(\{x_1,\dots,x_m\}) &= A(\{x_2,\dots,x_m\}) \cup \{ x_1, \operatorname{lca}(x_1,x_2) \} \\ &= A(\{x_3,\dots,x_m\}) \cup \{ x_1, \operatorname{lca}(x_1,x_2), x_2, \operatorname{lca}(x_2,x_3) \} \\ &\quad \vdots \\ &= \{ x_1, \operatorname{lca}(x_1,x_2), x_2, \dots, x_{m-1}, \operatorname{lca}(x_{m-1},x_m), x_m \}. \end{aligned} $$$

Таким образом, получаем требуемое представление.

Построение виртуального дерева

Дано множество вершин $$$X$$$. Виртуальное дерево $$$A(X)$$$ можно построить за время

$$$ O\left(|X| (\log |X| + \log N)\right), $$$

при этом необходимо выполнить предпосчет исходного дерева $$$T$$$ за $$$O(N \log N)$$$ для обеспечения быстрых запросов LCA.

Основные шаги построения:

  1. Предподсчет.
    С помощью обхода в глубину (DFS) или методов, таких как двоичные подъемы или HLD, вычисляем времена входа в вершины и LCA для любых двух вершин.

  2. Сортировка вершин.
    Сортируем вершины множества $$$X$$$ в порядке их появления при обходе DFS (используя время входа).

  3. Добавление LCA.
    Для последовательных вершин $$$x_i$$$ и $$$x_{i+1}$$$ вычисляем $$$\operatorname{lca}(x_i,x_{i+1})$$$. Объединение $$$X$$$ и этих LCA даёт множество вершин $$$A(X)$$$ (согласно свойству 2).

  4. Построение дерева.
    Получив множество вершин $$$A(X)$$$, отсортированных по порядку DFS, можно пройти по последовательности, используя стек, для восстановления структуры дерева. При обходе поддерживается стек, верхний элемент которого является текущей вершиной, а если новая вершина не является потомком вершины на вершине стека, производится «выталкивание» элементов до нахождения соответствующего предка, после чего производится соединение.

Такой алгоритм гарантирует, что виртуальное дерево содержит $$$O(|X|)$$$ вершин, что позволяет эффективно применять ДП.

Применение: подсчет поддеревьев в раскрашенном дереве

Условие задачи

Дано дерево $$$T$$$ с $$$N$$$ вершинами, пронумерованными от $$$1$$$ до $$$N$$$. Ребро $$$i$$$ соединяет вершины $$$u[i]$$$ и $$$v[i]$$$. Каждая вершина $$$i$$$ раскрашена в цвет $$$A[i]$$$.
Необходимо найти (по модулю 998244353) количество (непустых) подмножеств $$$S$$$ вершин дерева $$$T$$$, удовлетворяющих условию: - Индуцированный граф $$$G[S]$$$ является деревом. - Все вершины $$$G[S]$$$ со степенью 1 имеют один и тот же цвет.

Основная идея решения

Основная идея заключается в разбиении задачи по цветам. Для каждого цвета $$$c$$$ рассматривается множество вершин $$$X$$$, раскрашенных в $$$c$$$, и строится виртуальное дерево для $$$X$$$. Затем на полученном дереве выполняется ДП для подсчета допустимых поддеревьев, у которых все вершины со степенью 1 имеют цвет $$$c$$$. Итоговый ответ получается суммированием результатов по всем цветам.

Благодаря построению виртуального дерева, хотя исходное дерево содержит $$$N$$$ вершин, каждое виртуальное дерево для конкретного цвета имеет размер $$$O(|X|)$$$, что позволяет выполнять DP за приемлемое время.

Получаем, что суммарная сложность предпосчета и построения виртуальных деревьев не превышает $$$O(N \log N + \sum_{c \in C} |X_c| (\log |X_c| + \log N)) = O(N \log N)$$$, где $$$C$$$ это множество различных цветов вершин и $$$X_c$$$ это множество вершин с цветом $$$c$$$.

Реализация решения

Ниже приведён полный код на C++ с комментариями, описывающими основные компоненты алгоритма:

#include <bits/stdc++.h>
using namespace std;

const int MAXN = 2e5 + 10;
const int LOGN = 20;
const int MOD = 998244353;

vector<int> g[MAXN];         // Список смежности заданного графа
vector<int> ver[MAXN];       // Для каждого цвета ver[c] хранит вершины с цветом c
vector<int> aux_g[MAXN];     // Список смежности для виртуального дерева
int col[MAXN];               // Цвет каждой вершины
int tmr, n;                  // Глобальный счетчик времени и количество вершин
int up[LOGN][MAXN], dep[MAXN], tin[MAXN]; // Для подсчета LCA и времени входа

// DP массивы для динамического программирования на виртуальном дереве
int dp[MAXN][2], sum[MAXN];

// Функция предподсчета: DFS для вычисления tin, массива up и dep
void dfs_precalc(int v, int p) {
    tin[v] = ++tmr;
    up[0][v] = p;
    dep[v] = dep[p] + 1;
    for (int i = 1; i < LOGN; ++i)
        up[i][v] = up[i - 1][up[i - 1][v]];
    for (auto to : g[v])
        if (to != p)
            dfs_precalc(to, v);
}

// Функция для вычисления LCA с использованием двоичных подъемов
int getlca(int x, int y) {
    if (dep[x] < dep[y]) swap(x, y);
    for (int i = LOGN - 1; i >= 0; --i)
        if (dep[up[i][x]] >= dep[y])
            x = up[i][x];
    if (x == y) return x;
    for (int i = LOGN - 1; i >= 0; --i)
        if (up[i][x] != up[i][y]) {
            x = up[i][x];
            y = up[i][y];
        }
    return up[0][x];
}

// DFS по виртуальному дереву для выполнения DP.
// Параметр c — целевой цвет, для которого производится подсчет.

void dfs_calc(int v, int p, int c, int &ans) {
    dp[v][0] = dp[v][1] = 0;
    sum[v] = 0;
    for(auto to : aux_g[v]) {
        if(to == p) continue;
        dfs_calc(to, v, c, ans);
        // Переходы DP: объединение текущего состояния с результатом из поддерева.
        int nxt0 = (dp[v][0] + sum[to]) % MOD;
        int nxt1 = ((dp[v][0] + dp[v][1]) * 1ll * sum[to] % MOD + dp[v][1]) % MOD;
        dp[v][0] = nxt0;
        dp[v][1] = nxt1;
    }
    sum[v] = (dp[v][0] + dp[v][1]) % MOD;
    if(col[v] == c) {
        // Если вершина имеет целевой цвет, она может участвовать в корректном поддереве.
        sum[v] = (sum[v] + 1) % MOD;
        ans = (ans + sum[v]) % MOD;
    } else {
        ans = (ans + dp[v][1]) % MOD;
    }
}

// Функция для построения виртуального дерева для цвета c и выполнения DP.
void calc_aux(int c, int &ans) {
    auto p = ver[c];
    if (p.empty()) return;
    // Сортируем вершины по времени входа (tin) — порядку обхода DFS.
    sort(p.begin(), p.end(), [&](const int a, const int b) { return tin[a] < tin[b]; });
    vector<int> verstk = {1}; // Инициализируем стек с корнем дерева T (вершина 1).
    aux_g[1].clear();
    auto add = [&](int u, int v) {
        aux_g[u].push_back(v);
        aux_g[v].push_back(u);
    };
    // Обрабатываем каждую вершину из множества p, поддерживая стек для построения виртуального дерева.
    for (auto u : p) {
        if (u == 1) continue;
        int lca = getlca(u, verstk.back());
        if (lca != verstk.back()) {
            while (verstk.size() >= 2 && tin[lca] < tin[verstk[verstk.size() - 2]]) {
                add(verstk.back(), verstk[verstk.size() - 2]);
                verstk.pop_back();
            }
            if (verstk.size() >= 2 && tin[lca] != tin[verstk[verstk.size() - 2]]) {
                aux_g[lca].clear();
                add(verstk.back(), lca);
                verstk.back() = lca;
            } else {
                add(verstk.back(), lca);
                verstk.pop_back();
            }
        }
        aux_g[u].clear();
        verstk.push_back(u);
    }
    while (verstk.size() > 1) {
        add(verstk.back(), verstk[verstk.size() - 2]);
        verstk.pop_back();
    }
    // Выполняем DP по виртуальному дереву, начиная с корня 1.
    return dfs_calc(1, 0, c, ans);
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> n;
    // Считываем цвета и группируем вершины по цвету.
    for (int i = 1; i <= n; ++i) {
        cin >> col[i];
        ver[col[i]].push_back(i);
    }
    for (int i = 1; i < n; ++i) {
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs_precalc(1, 0);
    int ans = 0;
    // Для каждого возможного цвета обрабатываем соответствующее виртуальное дерево.
    for (int i = 1; i <= n; ++i)
        calc_aux(i, ans);
    cout << ans;
    return 0;
}

Объяснение реализации

  • Предподсчет.
    Функция dfs_precalc выполняет обход дерева $$$T$$$ в глубину, вычисляя время входа (tin), глубину и заполняя таблицу двоичных подъемов up для быстрых запросов LCA.

  • Вычисление LCA.
    Функция getlca реализует двоичные подъемы для быстрого нахождения наименьшего общего предка двух вершин.

  • Построение виртуального дерева.
    Функция calc_aux для каждого цвета $$$c$$$ берёт множество вершин ver[c], сортирует его по tin и с помощью стека строит виртуальное дерево. При этом для каждой пары последовательных вершин вычисляется LCA, что соответствует свойству 2.

  • Динамическое программирование.
    Функция dfs_calc выполняет обход виртуального дерева и объединяет результаты поддеревьев согласно переходам DP. Состояния DP рассчитаны таким образом, что учитывается вклад вершины, если она имеет целевой цвет, и производится подсчет только тех поддеревьев, где все листья имеют один цвет.

  • Сбор результата.
    Функция main считывает входные данные, выполняет предподсчет, и затем суммирует результаты для каждого цвета, выводя итоговый ответ по модулю 998244353.

Заключение

В данной статье мы подробно рассмотрели понятие виртуального дерева, доказали его основные свойства, описали эффективный алгоритм построения и продемонстрировали применение этой техники для решения задачи на ДП в деревьях. Преимущество данного подхода состоит в том, что он позволяет свести задачу к работе с подструктурой размером $$$O(|X|)$$$, что существенно ускоряет вычисления даже при большом размере исходного дерева.

Статья основана на разборе задачи с Atcoder (ABC340G) от Nyaan.

  • Проголосовать: нравится
  • +29
  • Проголосовать: не нравится

»
7 часов назад, # |
  Проголосовать: нравится 0 Проголосовать: не нравится

Auto comment: topic has been updated by spike1236 (previous revision, new revision, compare).

→ Ответить
»
7 часов назад, # |
  Проголосовать: нравится 0 Проголосовать: не нравится

Автокомментарий: текст был обновлен пользователем spike1236 (предыдущая версия, новая версия, сравнить).

→ Ответить
»
2 часа назад, # |
  Проголосовать: нравится 0 Проголосовать: не нравится

thanks!

→ Ответить