Implementação de Heavy-Light Decomposition e Segment Tree com Prioridade de Lazy Tags

Resolver problemas que exigem consultas e atualizações em caminhos de árvores frequentemetne requer a conversão de pesos de arestas para pesos de vértices, utilizando a técnica de Heavy-Light Decomposition (HLD), também conhecida como Decomposição em Cadeias Pesadas. O núcleo da complexidade deste tipo de problema reside na implementação correta da Segment Tree (Árvore de Segmentos), especialmente quando é necessário gerenciar múltiplas operações de atualização de forma simultânea.

Gerenciamento de Conflitos entre Lazy Tags

Para suportar as operações de modificação de caminhos, a estrutura de dados deve manter a soma, o valor mínimo e o valor máximo dos segmentos. O desafio algorítmico surge ao introduzir dois tipos de lazy tags (marcações preguiçosas): uma para atribuição absoluta de valor (cobertura ou cover) e outra para incremento cumulativo (adição ou add).

A coexistência de operações de sobrescrita e adição exige uma definição estrita de precedência. Quando uma operação de atribuição é aplicada a um nó da árvore de segmentos, qualquer incremento penednte nesse mesmo nó torna-se obsoleto e deve ser imediatamente descartado. Portanto, a regra de ouro é: a marcação de atribuição possui prioridade absoluta e anula a marcação de adição existente.

Considerando que os pesos das arestas na árvore são estritamente positivos, podemos utilizar -1 como um valor sentinela para indicar a ausência de uma marcação de atribuição. Durante a propagação das marcas para os nós filhos (operação de pushdown), a ordem de execução é crítica:

  1. Propagação de Atribuição: Verifique se existe uma marca de atribuição (lazy_set != -1). Se existir, aplique o valor aos filhos, atualize suas somas, mínimos e máximos. Em seguida, sobrescreva as marcas de atribuição dos filhos e force as marcas de adição dos filhos para zero.
  2. Propagação de Adição: Verifique se existe um incremento pendente (lazy_add != 0). Se existir, acumule o valor nas somas, mínimos e máximos dos filhos, e adicione-o às marcas de adição já existentes neles.

Implementação da Estrutura de Dados

Abaixo está a implementação refatorada do algoritmo. As variáveis foram renomeadas para maior clareza semântica e a lógica de manipulação de bits foi substituída por operações aritméticas diretas. A estrutura engloba a leitura da árvore, a linearização via HLD e a Segment Tree com tratamento de prioridades de lazy tags.

#include <iostream>
#include <vector>
#include <algorithm>
#include <string>

using namespace std;

const int MAX_NODES = 100005;
const int INF = 1e9 + 7;

struct Edge {
    int to, weight;
};

struct QueryEdge {
    int u, v;
} original_edges[MAX_NODES];

vector<Edge> adj[MAX_NODES];

// Heavy-Light Decomposition Arrays
int parent_node[MAX_NODES], depth[MAX_NODES], subtree_size[MAX_NODES];
int heavy_child[MAX_NODES], chain_top[MAX_NODES], dfn[MAX_NODES];
int node_weight[MAX_NODES], edge_to_node_weight[MAX_NODES];
int dfn_counter = 0;

// Segment Tree Node
struct SegNode {
    int left, right;
    long long sum;
    int min_val, max_val;
    int lazy_set, lazy_add;
} seg_tree[MAX_NODES << 2];

inline int get_length(int node) {
    return seg_tree[node].right - seg_tree[node].left + 1;
}

void pushup(int node) {
    int lc = node << 1, rc = node << 1 | 1;
    seg_tree[node].sum = seg_tree[lc].sum + seg_tree[rc].sum;
    seg_tree[node].min_val = min(seg_tree[lc].min_val, seg_tree[rc].min_val);
    seg_tree[node].max_val = max(seg_tree[lc].max_val, seg_tree[rc].max_val);
}

void build_tree(int node, int l, int r) {
    seg_tree[node].left = l;
    seg_tree[node].right = r;
    seg_tree[node].lazy_set = -1;
    seg_tree[node].lazy_add = 0;
    
    if (l == r) {
        seg_tree[node].min_val = seg_tree[node].max_val = seg_tree[node].sum = node_weight[l];
        return;
    }
    int mid = l + (r - l) / 2;
    build_tree(node << 1, l, mid);
    build_tree(node << 1 | 1, mid + 1, r);
    pushup(node);
}

void propagate(int node) {
    int lc = node << 1, rc = node << 1 | 1;
    int len_lc = get_length(lc), len_rc = get_length(rc);

    // Priority 1: Assignment overrides addition
    if (seg_tree[node].lazy_set != -1) {
        int val = seg_tree[node].lazy_set;
        
        seg_tree[lc].sum = 1LL * len_lc * val;
        seg_tree[lc].min_val = seg_tree[lc].max_val = val;
        seg_tree[lc].lazy_set = val;
        seg_tree[lc].lazy_add = 0; // Clear addition tag
        
        seg_tree[rc].sum = 1LL * len_rc * val;
        seg_tree[rc].min_val = seg_tree[rc].max_val = val;
        seg_tree[rc].lazy_set = val;
        seg_tree[rc].lazy_add = 0; // Clear addition tag
        
        seg_tree[node].lazy_set = -1;
    }
    
    // Priority 2: Accumulate addition
    if (seg_tree[node].lazy_add != 0) {
        int val = seg_tree[node].lazy_add;
        
        seg_tree[lc].sum += 1LL * len_lc * val;
        seg_tree[lc].min_val += val;
        seg_tree[lc].max_val += val;
        seg_tree[lc].lazy_add += val;
        
        seg_tree[rc].sum += 1LL * len_rc * val;
        seg_tree[rc].min_val += val;
        seg_tree[rc].max_val += val;
        seg_tree[rc].lazy_add += val;
        
        seg_tree[node].lazy_add = 0;
    }
}

void update_set(int node, int l, int r, int val) {
    if (l <= seg_tree[node].left && seg_tree[node].right <= r) {
        seg_tree[node].sum = 1LL * get_length(node) * val;
        seg_tree[node].min_val = seg_tree[node].max_val = val;
        seg_tree[node].lazy_set = val;
        seg_tree[node].lazy_add = 0;
        return;
    }
    propagate(node);
    int mid = seg_tree[node].left + (seg_tree[node].right - seg_tree[node].left) / 2;
    if (l <= mid) update_set(node << 1, l, r, val);
    if (r > mid) update_set(node << 1 | 1, l, r, val);
    pushup(node);
}

void update_add(int node, int l, int r, int val) {
    if (l <= seg_tree[node].left && seg_tree[node].right <= r) {
        seg_tree[node].sum += 1LL * get_length(node) * val;
        seg_tree[node].min_val += val;
        seg_tree[node].max_val += val;
        seg_tree[node].lazy_add += val;
        return;
    }
    propagate(node);
    int mid = seg_tree[node].left + (seg_tree[node].right - seg_tree[node].left) / 2;
    if (l <= mid) update_add(node << 1, l, r, val);
    if (r > mid) update_add(node << 1 | 1, l, r, val);
    pushup(node);
}

int query_max(int node, int l, int r) {
    if (l <= seg_tree[node].left && seg_tree[node].right <= r) {
        return seg_tree[node].max_val;
    }
    propagate(node);
    int mid = seg_tree[node].left + (seg_tree[node].right - seg_tree[node].left) / 2;
    int res = -INF;
    if (l <= mid) res = max(res, query_max(node << 1, l, r));
    if (r > mid) res = max(res, query_max(node << 1 | 1, l, r));
    return res;
}

// HLD DFS 1: Compute subtree sizes and heavy children
void dfs1(int u, int p, int d) {
    parent_node[u] = p;
    depth[u] = d;
    subtree_size[u] = 1;
    heavy_child[u] = 0;
    
    for (const auto& edge : adj[u]) {
        int v = edge.to;
        if (v == p) continue;
        edge_to_node_weight[v] = edge.weight;
        dfs1(v, u, d + 1);
        subtree_size[u] += subtree_size[v];
        if (subtree_size[v] > subtree_size[heavy_child[u]]) {
            heavy_child[u] = v;
        }
    }
}

// HLD DFS 2: Map nodes to segment tree indices
void dfs2(int u, int top) {
    dfn[u] = ++dfn_counter;
    node_weight[dfn_counter] = edge_to_node_weight[u];
    chain_top[u] = top;
    
    if (!heavy_child[u]) return;
    dfs2(heavy_child[u], top);
    
    for (const auto& edge : adj[u]) {
        int v = edge.to;
        if (v == parent_node[u] || v == heavy_child[u]) continue;
        dfs2(v, v);
    }
}

// Path operations applying edge-to-node conversion logic
void path_update_set(int u, int v, int val) {
    while (chain_top[u] != chain_top[v]) {
        if (depth[chain_top[u]] < depth[chain_top[v]]) swap(u, v);
        update_set(1, dfn[chain_top[u]], dfn[u], val);
        u = parent_node[chain_top[u]];
    }
    if (depth[u] > depth[v]) swap(u, v);
    if (u != v) update_set(1, dfn[u] + 1, dfn[v], val); // Skip LCA
}

void path_update_add(int u, int v, int val) {
    while (chain_top[u] != chain_top[v]) {
        if (depth[chain_top[u]] < depth[chain_top[v]]) swap(u, v);
        update_add(1, dfn[chain_top[u]], dfn[u], val);
        u = parent_node[chain_top[u]];
    }
    if (depth[u] > depth[v]) swap(u, v);
    if (u != v) update_add(1, dfn[u] + 1, dfn[v], val); // Skip LCA
}

int path_query_max(int u, int v) {
    int res = -INF;
    while (chain_top[u] != chain_top[v]) {
        if (depth[chain_top[u]] < depth[chain_top[v]]) swap(u, v);
        res = max(res, query_max(1, dfn[chain_top[u]], dfn[u]));
        u = parent_node[chain_top[u]];
    }
    if (depth[u] > depth[v]) swap(u, v);
    if (u != v) res = max(res, query_max(1, dfn[u] + 1, dfn[v])); // Skip LCA
    return res;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    
    int n;
    if (!(cin >> n)) return 0;
    
    for (int i = 1; i < n; ++i) {
        int u, v, w;
        cin >> u >> v >> w;
        original_edges[i] = {u, v};
        adj[u].push_back({v, w});
        adj[v].push_back({u, w});
    }
    
    dfs1(1, 0, 1);
    dfs2(1, 1);
    build_tree(1, 1, n);
    
    string op;
    while (cin >> op && op != "Stop") {
        if (op == "Change") {
            int k, w;
            cin >> k >> w;
            path_update_set(original_edges[k].u, original_edges[k].v, w);
        } else if (op == "Cover") {
            int u, v, w;
            cin >> u >> v >> w;
            path_update_set(u, v, w);
        } else if (op == "Add") {
            int u, v, w;
            cin >> u >> v >> w;
            path_update_add(u, v, w);
        } else if (op == "Max") {
            int u, v;
            cin >> u >> v;
            if (u == v) cout << 0 << "\n";
            else cout << path_query_max(u, v) << "\n";
        }
    }
    return 0;
}

Gerador de Casos de Teste

Para validar a corretude da manipulação das lazy tags, é fundamental realizar testes de estresse (stress testing) contra uma implementação ingênua. O código abaixo utiliza a biblioteca <random> do C++11 para gerar topologias de árvores conectadas e uma sequência aleatória de operações válidas de forma mais robusta que o gerador padrão rand().

#include <iostream>
#include <fstream>
#include <random>
#include <string>

using namespace std;

int main() {
    ofstream out_file("stress_test.in");
    if (!out_file.is_open()) return 1;

    // Seed the Mersenne Twister pseudo-random generator
    mt19937 rng(42); 
    
    int num_nodes = 100;
    int num_queries = 1000;
    
    out_file << num_nodes << "\n";
    
    // Generate a valid connected tree topology
    uniform_int_distribution<int> weight_dist(1, 1000);
    for (int i = 2; i <= num_nodes; ++i) {
        uniform_int_distribution<int> parent_dist(1, i - 1);
        out_file << parent_dist(rng) << " " << i << " " << weight_dist(rng) << "\n";
    }
    
    // Query generation parameters
    string operations[] = {"Max", "Add", "Cover", "Change"};
    uniform_int_distribution<int> op_dist(0, 3);
    uniform_int_distribution<int> node_dist(1, num_nodes);
    uniform_int_distribution<int> edge_dist(1, num_nodes - 1);
    
    for (int i = 0; i < num_queries; ++i) {
        int current_op = op_dist(rng);
        
        if (current_op == 0) { // Max
            out_file << operations[current_op] << " " 
                     << node_dist(rng) << " " 
                     << node_dist(rng) << "\n";
        } else if (current_op == 1 || current_op == 2) { // Add, Cover
            out_file << operations[current_op] << " " 
                     << node_dist(rng) << " " 
                     << node_dist(rng) << " " 
                     << weight_dist(rng) << "\n";
        } else { // Change
            out_file << operations[current_op] << " " 
                     << edge_dist(rng) << " " 
                     << weight_dist(rng) << "\n";
        }
    }
    
    out_file << "Stop\n";
    out_file.close();
    
    return 0;
}

Tags: C++ Heavy-Light Decomposition segment tree Lazy Propagation Graph Algorithms

Publicado em 7-5 00:57