kz-maxx

DP на дереве

Динамическое программирование по поддеревьям через DFS

Что такое DP на дереве?

DP на дереве — это динамическое программирование, где состояние считается для каждой вершины дерева.

Обычно мы запускаем DFS и считаем ответ для поддерева каждой вершины.

сначала считаем ответы для детей
потом через них считаем ответ для родителя
Главное правило: чтобы посчитать dp[v], нужно сначала посчитать dp для всех детей v.

Что такое дерево?

Дерево — это связный граф без циклов.

Если в дереве n вершин, то ребер всегда n - 1.

дерево = связный граф + нет циклов
количество ребер = n - 1

Корень дерева

Чтобы делать DP на дереве, дерево часто подвешивают за корень.

Обычно выбирают вершину 1 как корень:

dfs(1, 0);

Тогда у каждой вершины появляются:

  • parent — родитель
  • children — дети
  • subtree — поддерево вершины

Базовый DFS на дереве

В дереве нет циклов, но если хранить неориентированные ребра, нужно не возвращаться в родителя.

void dfs(int v, int p) {
    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs(to, v);
    }
}

Здесь:

  • v — текущая вершина
  • p — родитель вершины v
  • to — сосед вершины v

Главная схема DP на дереве

void dfs(int v, int p) {
    // базовое значение dp[v]

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs(to, v);

        // обновляем dp[v] через dp[to]
    }
}
Это самая важная идея: сначала dfs(to), потом используем dp[to].

Задача 1: размер поддерева

Для каждой вершины нужно найти размер ее поддерева.

Состояние:

dp[v] = размер поддерева вершины v

Сначала сама вершина уже дает размер 1:

dp[v] = 1;

Потом добавляем размеры всех детей:

dp[v] += dp[to];

Код: размер поддерева

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

const int N = 100010;

int n;
vector<int> g[N];
int dp[N];

void dfs(int v, int p) {
    dp[v] = 1;

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs(to, v);
        dp[v] += dp[to];
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;

    for (int i = 1; i <= n - 1; i++) {
        int a, b;
        cin >> a >> b;

        g[a].push_back(b);
        g[b].push_back(a);
    }

    dfs(1, 0);

    for (int i = 1; i <= n; i++) {
        cout << dp[i] << ' ';
    }

    return 0;
}

Пример для размера поддерева

Input

5
1 2
1 3
3 4
3 5

Если корень 1, дерево выглядит так:

1
├── 2
└── 3
    ├── 4
    └── 5

Размеры поддеревьев:

dp[1] = 5
dp[2] = 1
dp[3] = 3
dp[4] = 1
dp[5] = 1

Задача 2: максимальная сумма независимых вершин

Дано дерево. У каждой вершины есть вес w[v]. Нужно выбрать вершины с максимальной суммой весов так, чтобы никакие две выбранные вершины не были соседями.

Это называется Maximum Independent Set on Tree.

если выбрали вершину v
то нельзя выбрать ее детей

Состояния для независимого множества

Для каждой вершины будем хранить два значения:

dp[v][0]
dp[v][1]
  • dp[v][0] — лучший ответ в поддереве v, если v не выбрана
  • dp[v][1] — лучший ответ в поддереве v, если v выбрана

Переход

Если вершина v выбрана, то детей выбирать нельзя:

dp[v][1] += dp[to][0];

Если вершина v не выбрана, то каждый ребенок может быть выбран или не выбран:

dp[v][0] += max(dp[to][0], dp[to][1]);

Начальное значение:

dp[v][1] = w[v];
dp[v][0] = 0;

Код: максимальная сумма независимых вершин

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

const int N = 100010;

int n;
vector<int> g[N];
long long w[N];
long long dp[N][2];

void dfs(int v, int p) {
    dp[v][0] = 0;
    dp[v][1] = w[v];

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs(to, v);

        dp[v][0] += max(dp[to][0], dp[to][1]);
        dp[v][1] += dp[to][0];
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;

    for (int i = 1; i <= n; i++) {
        cin >> w[i];
    }

    for (int i = 1; i <= n - 1; i++) {
        int a, b;
        cin >> a >> b;

        g[a].push_back(b);
        g[b].push_back(a);
    }

    dfs(1, 0);

    cout << max(dp[1][0], dp[1][1]);

    return 0;
}

Пример

Input

5
10 1 5 7 8
1 2
1 3
3 4
3 5

Можно выбрать вершины:

1, 4, 5
сумма = 10 + 7 + 8 = 25

Output

25

Задача 3: минимальное вершинное покрытие дерева

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

Это называется Minimum Vertex Cover on Tree.

для каждого ребра a--b
нужно выбрать a или b

Состояния для vertex cover

dp[v][0]
dp[v][1]
  • dp[v][0] — минимум в поддереве v, если v не выбрана
  • dp[v][1] — минимум в поддереве v, если v выбрана

Если v не выбрана, то все дети обязательно должны быть выбраны.

dp[v][0] += dp[to][1];

Если v выбрана, то ребенок может быть выбран или не выбран.

dp[v][1] += min(dp[to][0], dp[to][1]);

Код: минимальное вершинное покрытие

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

const int N = 100010;

int n;
vector<int> g[N];
int dp[N][2];

void dfs(int v, int p) {
    dp[v][0] = 0;
    dp[v][1] = 1;

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs(to, v);

        dp[v][0] += dp[to][1];
        dp[v][1] += min(dp[to][0], dp[to][1]);
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;

    for (int i = 1; i <= n - 1; i++) {
        int a, b;
        cin >> a >> b;

        g[a].push_back(b);
        g[b].push_back(a);
    }

    dfs(1, 0);

    cout << min(dp[1][0], dp[1][1]);

    return 0;
}

Задача 4: высота поддерева

Нужно для каждой вершины найти высоту ее поддерева.

Состояние:

dp[v] = высота поддерева v

Если у вершины нет детей, ее высота равна 1.

Иначе:

dp[v] = max(dp[to] + 1)

Код: высота поддерева

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

const int N = 100010;

int n;
vector<int> g[N];
int dp[N];

void dfs(int v, int p) {
    dp[v] = 1;

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs(to, v);

        dp[v] = max(dp[v], dp[to] + 1);
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;

    for (int i = 1; i <= n - 1; i++) {
        int a, b;
        cin >> a >> b;

        g[a].push_back(b);
        g[b].push_back(a);
    }

    dfs(1, 0);

    for (int i = 1; i <= n; i++) {
        cout << dp[i] << ' ';
    }

    return 0;
}

Задача 5: диаметр дерева через DP

Диаметр дерева — это максимальное расстояние между двумя вершинами.

Его можно искать через DP на дереве.

Для каждой вершины храним две самые большие высоты среди ее детей.

путь может идти через вершину v
из одного поддерева в другое поддерево

Идея диаметра

Пусть best1 и best2 — две максимальные глубины от детей.

Тогда возможный диаметр через v:

best1 + best2

А высота для родителя:

dp[v] = best1 + 1;

Код: диаметр дерева

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

const int N = 100010;

int n;
vector<int> g[N];
int dp[N];
int ans = 0;

void dfs(int v, int p) {
    int best1 = 0;
    int best2 = 0;

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs(to, v);

        int cur = dp[to] + 1;

        if (cur > best1) {
            best2 = best1;
            best1 = cur;
        } else if (cur > best2) {
            best2 = cur;
        }
    }

    dp[v] = best1;
    ans = max(ans, best1 + best2);
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;

    for (int i = 1; i <= n - 1; i++) {
        int a, b;
        cin >> a >> b;

        g[a].push_back(b);
        g[b].push_back(a);
    }

    dfs(1, 0);

    cout << ans;

    return 0;
}
Этот код считает диаметр в количестве ребер. Если нужно количество вершин, ответ будет ans + 1.

Рекурсивный метод в DP на дереве

Почти все DP на дереве пишется рекурсивно.

Мы запускаем DFS:

dfs(v, parent)

И внутри DFS считаем dp[v].

Это удобно, потому что дерево само задает порядок:

листы считаются первыми
их родители после них
корень считается последним

Как конвертировать в итеративный метод?

Иногда рекурсию нужно заменить на итеративный DFS, чтобы избежать stack overflow.

Идея:

1. Сначала получаем порядок обхода вершин
2. Потом идем по этому порядку с конца
3. Тогда дети будут обработаны раньше родителей

Итеративный порядок для DP

vector<int> order;
vector<int> parent(n + 1, 0);

stack<int> st;
st.push(1);
parent[1] = -1;

while (!st.empty()) {
    int v = st.top();
    st.pop();

    order.push_back(v);

    for (int to : g[v]) {
        if (to == parent[v]) {
            continue;
        }

        parent[to] = v;
        st.push(to);
    }
}

reverse(order.begin(), order.end());

После reverse сначала идут дети, потом родители. Это как DFS postorder.

Итеративный DP для размера поддерева

for (int v : order) {
    dp[v] = 1;

    for (int to : g[v]) {
        if (to == parent[v]) {
            continue;
        }

        dp[v] += dp[to];
    }
}

Полная итеративная реализация

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

const int N = 100010;

int n;
vector<int> g[N];
int parent[N];
int dp[N];

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;

    for (int i = 1; i <= n - 1; i++) {
        int a, b;
        cin >> a >> b;

        g[a].push_back(b);
        g[b].push_back(a);
    }

    vector<int> order;
    stack<int> st;

    st.push(1);
    parent[1] = -1;

    while (!st.empty()) {
        int v = st.top();
        st.pop();

        order.push_back(v);

        for (int to : g[v]) {
            if (to == parent[v]) {
                continue;
            }

            parent[to] = v;
            st.push(to);
        }
    }

    reverse(order.begin(), order.end());

    for (int v : order) {
        dp[v] = 1;

        for (int to : g[v]) {
            if (to == parent[v]) {
                continue;
            }

            dp[v] += dp[to];
        }
    }

    for (int i = 1; i <= n; i++) {
        cout << dp[i] << ' ';
    }

    return 0;
}

Как придумывать DP на дереве?

Обычно нужно ответить на вопросы:

1. Что означает dp[v]?
2. Нужны ли состояния 0/1?
3. Что происходит, если вершина v выбрана?
4. Что происходит, если вершина v не выбрана?
5. Как объединить ответы детей?
6. Где лежит ответ: в dp[1] или max/min из dp[1]?

Типичные состояния

Состояние Что означает
dp[v] ответ для поддерева v
dp[v][0] ответ, если v не выбрана
dp[v][1] ответ, если v выбрана
subtree[v] размер поддерева v
height[v] высота поддерева v

Типичные ошибки

  • Не передать родителя в DFS
  • Зайти обратно в родителя и получить бесконечную рекурсию
  • Обновлять dp[v] до dfs(to)
  • Забыть базовое значение dp[v]
  • Перепутать dp[v][0] и dp[v][1]
  • Взять ответ только dp[1][1], хотя нужен max или min
  • Получить stack overflow на большом дереве

Асимптотика

Задача Время Память
Размер поддерева O(n) O(n)
Independent Set O(n) O(n)
Vertex Cover O(n) O(n)
Диаметр O(n) O(n)
Обычно DP на дереве работает за O(n), потому что DFS проходит каждую вершину и каждое ребро один раз.

Короткий шаблон DP на дереве

void dfs(int v, int p) {
    // base dp[v]

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs(to, v);

        // update dp[v] using dp[to]
    }
}

Короткий шаблон dp[v][0/1]

void dfs(int v, int p) {
    dp[v][0] = base0;
    dp[v][1] = base1;

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs(to, v);

        dp[v][0] += transition_when_v_not_taken;
        dp[v][1] += transition_when_v_taken;
    }
}