Colorir as arestas de uma árvore

Esta é uma tarefa bastante comum. Dada uma árvore $G$ com $N$ vértices. Existem dois tipos de consultas: a primeira é colorir uma aresta, a segunda é o número de arestas coloridas entre dois vértices.

Aqui, descreveremos uma solução bastante simples (usando uma árvore segmentária) que responderá a cada consulta em $O(\log N)$. A etapa de pré-processamento levará $O(N)$.

Algoritmo

Primeiro, precisamos encontrar o LCA(menor ancestral comum) para reduzir cada consulta do segundo tipo(o número de arestas coloridas entre dois vértices) $(i,j)$ em duas consultas: $(l,i)$ e $(l,j)$; onde $l$ é o LCA de $i$ e $j$. A resposta da consulta $(i,j)$ será a soma das duas subconsultas. Ambas as consultas têm uma estrutura especial, o primeiro vértice é um ancestral do segundo. No restante do artigo, falaremos apenas sobre esse tipo especial de consulta.

Começaremos descrevendo a etapa de pré-processamento. Execute uma DFS a partir da raiz da árvore e registre o tour de Euler dessa DFS (cada vértice é adicionado à lista quando a pesquisa o visita primeiro e toda vez que retornamos de um de seus filhos). A mesma técnica pode ser usada no pré-processamento do LCA.

Essa lista conterá cada aresta (no sentido de que se $i$ e $j$ forem as extremidades das arestas, haverá um local na lista em que $i$ e $j$ são vizinhos na lista), e aparece exatamente duas vezes: na direção direta (de $i$ a $j$, onde o vértice $i$ está mais próximo da raiz que o vértice $j$) e na direção oposta (de $j$ a $i$).

Vamos construir duas listas para essas arestas. A primeiro armazena a cor de todas as arestas na direção direta, e a segunda a cor de todas as arestas na direção oposta. Usaremos $1$ se a aresta for colorida, e $0$ caso contrário. Sobre essas duas listas construiremos para cada uma, uma árvore de segmentos (para a soma), vamos chamá-las de $T1$ e $T2$.

Vamos responder a uma consulta da forma "$(i,j)$", onde $i$ é o ancestral de $j$. Precisamos determinar quantas arestas são pintadas/coloridas no caminho entre $i$ e $j$. Vamos encontrar $i$ e $j$ no tour de Euler pela primeira vez, sejam as posições $p$ e $q$ (isso pode ser feito em $O(1)$ se calcularmos essas posições com antecedência durante pré-processamento). Assim, a resposta para a consulta é a soma $T1[p..q-1]$ menos a soma $T2[p..q-1]$.

Considere o segmento $[p;q]$ no tour de Euler. Ele contém todas as arestas que precisamos do caminho de $i$ a $j$, mas também contém um conjunto de arestas que se encontram em outros caminhos de $i$. No entanto, há uma grande diferença entre as arestas que precisamos e as demais: as arestas que precisamos serão listadas apenas uma vez na direção direta, e todas as outras arestas aparecerão duas vezes: nas duas direções(direta e oposta). Assim, a diferença $T1[p..q-1] - T2[p..q-1]$ nos dará a resposta correta ("menos um" é necessário porque, caso contrário, capturaremos uma aresta extra saindo do vértice $j$). A consulta de soma na árvore de segmentos é executada em $O(\log N)$.

Responder o primeiro tipo de consulta (colorir uma aresta) é ainda mais fácil - precisamos apenas atualizar $T1$ e $T2$, ou seja, realizar uma única atualização do elemento que corresponde à nossa aresta (encontrar a aresta na lista, novamente, é possível em $O(1)$, se você executar esta pesquisa durante o pré-processamento). Uma única modificação na árvore de segmentos é realizada em $O(\log N)$.

Implementação

Aqui está a implementação completa da solução, incluindo a computação do LCA:

const int INF = 1000 * 1000 * 1000;

typedef vector<vector<int>> graph;

vector<int> dfs_list;
vector<int> edges_list;
vector<int> h;

void dfs(int v, const graph& g, const graph& edge_ids, int cur_h = 1) {
    h[v] = cur_h;
    dfs_list.push_back(v);
    for (size_t i = 0; i < g[v].size(); ++i) {
        if (h[g[v][i]] == -1) {
            edges_list.push_back(edge_ids[v][i]);
            dfs(g[v][i], g, edge_ids, cur_h + 1);
            edges_list.push_back(edge_ids[v][i]);
            dfs_list.push_back(v);
        }
    }
}

vector<int> lca_tree;
vector<int> first;

void lca_tree_build(int i, int l, int r) {
    if (l == r) {
        lca_tree[i] = dfs_list[l];
    } else {
        int m = (l + r) >> 1;
        lca_tree_build(i + i, l, m);
        lca_tree_build(i + i + 1, m + 1, r);
        int lt = lca_tree[i + i], rt = lca_tree[i + i + 1];
        lca_tree[i] = h[lt] < h[rt] ? lt : rt;
    }
}

void lca_prepare(int n) {
    lca_tree.assign(dfs_list.size() * 8, -1);
    lca_tree_build(1, 0, (int)dfs_list.size() - 1);

    first.assign(n, -1);
    for (int i = 0; i < (int)dfs_list.size(); ++i) {
        int v = dfs_list[i];
        if (first[v] == -1)
            first[v] = i;
    }
}

int lca_tree_query(int i, int tl, int tr, int l, int r) {
    if (tl == l && tr == r)
        return lca_tree[i];
    int m = (tl + tr) >> 1;
    if (r <= m)
        return lca_tree_query(i + i, tl, m, l, r);
    if (l > m)
        return lca_tree_query(i + i + 1, m + 1, tr, l, r);
    int lt = lca_tree_query(i + i, tl, m, l, m);
    int rt = lca_tree_query(i + i + 1, m + 1, tr, m + 1, r);
    return h[lt] < h[rt] ? lt : rt;
}

int lca(int a, int b) {
    if (first[a] > first[b])
        swap(a, b);
    return lca_tree_query(1, 0, (int)dfs_list.size() - 1, first[a], first[b]);
}

vector<int> first1, first2;
vector<char> edge_used;
vector<int> tree1, tree2;

void query_prepare(int n) {
    first1.resize(n - 1, -1);
    first2.resize(n - 1, -1);
    for (int i = 0; i < (int)edges_list.size(); ++i) {
        int j = edges_list[i];
        if (first1[j] == -1)
            first1[j] = i;
        else
            first2[j] = i;
    }

    edge_used.resize(n - 1);
    tree1.resize(edges_list.size() * 8);
    tree2.resize(edges_list.size() * 8);
}

void sum_tree_update(vector<int>& tree, int i, int l, int r, int j, int delta) {
    tree[i] += delta;
    if (l < r) {
        int m = (l + r) >> 1;
        if (j <= m)
            sum_tree_update(tree, i + i, l, m, j, delta);
        else
            sum_tree_update(tree, i + i + 1, m + 1, r, j, delta);
    }
}

int sum_tree_query(const vector<int>& tree, int i, int tl, int tr, int l, int r) {
    if (l > r || tl > tr)
        return 0;
    if (tl == l && tr == r)
        return tree[i];
    int m = (tl + tr) >> 1;
    if (r <= m)
        return sum_tree_query(tree, i + i, tl, m, l, r);
    if (l > m)
        return sum_tree_query(tree, i + i + 1, m + 1, tr, l, r);
    return sum_tree_query(tree, i + i, tl, m, l, m) +
           sum_tree_query(tree, i + i + 1, m + 1, tr, m + 1, r);
}

int query(int v1, int v2) {
    return sum_tree_query(tree1, 1, 0, (int)edges_list.size() - 1, first[v1], first[v2] - 1) -
           sum_tree_query(tree2, 1, 0, (int)edges_list.size() - 1, first[v1], first[v2] - 1);
}

int main() {
    // lendo o grafo
    int n;
    scanf("%d", &n);
    graph g(n), edge_ids(n);
    for (int i = 0; i < n - 1; ++i) {
        int v1, v2;
        scanf("%d%d", &v1, &v2);
        --v1, --v2;
        g[v1].push_back(v2);
        g[v2].push_back(v1);
        edge_ids[v1].push_back(i);
        edge_ids[v2].push_back(i);
    }

    h.assign(n, -1);
    dfs(0, g, edge_ids);
    lca_prepare(n);
    query_prepare(n);

    for (;;) {
        if () {
            // pedido para colorir aresta x;
            // if start = true, então a aresta é colorida, caso contrário a coloração/pintura
            // é removida
            edge_used[x] = start;
            sum_tree_update(tree1, 1, 0, (int)edges_list.size() - 1, first1[x],
                            start ? 1 : -1);
            sum_tree_update(tree2, 1, 0, (int)edges_list.size() - 1, first2[x],
                            start ? 1 : -1);
        } else {
            // consulta o número de arestas coloridas no caminho entre v1 e v2
            int l = lca(v1, v2);
            int result = query(l, v1) + query(l, v2);
            // resultado - a resposta para o pedido
        }
    }
}