Pintar las aristas del árbol

Esta es una tarea bastante común. Se da un árbol $G$ con $N$ vértices. Hay dos tipos de consultas: la primera es pintar una arista, la segunda es consultar el número de aristas pintadas entre dos vértices.

Aquí describiremos una solución bastante simple (usando un Árbol de Segmentos) que responderá cada consulta en tiempo $O(\log N)$ . El paso de preprocesamiento tomará tiempo $O(N)$ .

Algoritmo

Primero, necesitamos encontrar el LCA para reducir cada consulta del segundo tipo $(i,j)$ a dos consultas $(l,i)$ y $(l,j)$ , donde $l$ es el LCA de $i$ y $j$ . La respuesta de la consulta $(i,j)$ será la suma de ambas subconsultas. Ambas consultas tienen una estructura especial: el primer vértice es un ancestro del segundo. Para el resto del artículo solo hablaremos de este tipo especial de consultas.

Empezaremos describiendo el paso de preprocesamiento. Ejecutamos una búsqueda en profundidad desde la raíz del árbol y registramos el tour de Euler de esta búsqueda en profundidad (cada vértice se agrega a la lista cuando la búsqueda lo visita por primera vez y cada vez que regresamos de uno de sus hijos). La misma técnica se puede usar en el preprocesamiento del LCA.

Esta lista contendrá cada arista (en el sentido de que si $i$ y $j$ son los extremos de la arista, entonces habrá un lugar en la lista donde $i$ y $j$ son vecinos en la lista), y aparecerá exactamente dos veces: en la dirección de ida (de $i$ a $j$ , donde el vértice $i$ está más cerca de la raíz que el vértice $j$ ) y en la dirección opuesta (de $j$ a $i$ ).

Construiremos dos listas para estas aristas. La primera guardará el color de todas las aristas en la dirección de ida, y la segunda el color de todas las aristas en la dirección opuesta. Usaremos $1$ si la arista está pintada, y $0$ en caso contrario. Sobre estas dos listas construiremos un Árbol de Segmentos para cada una (para suma con una sola modificación); llamémoslos $T1$ y $T2$ .

Respondamos una consulta de la forma $(i,j)$ , donde $i$ es el ancestro de $j$ . Necesitamos determinar cuántas aristas están pintadas en el camino entre $i$ y $j$ . Encontremos $i$ y $j$ en el tour de Euler por primera vez; sean las posiciones $p$ y $q$ (esto se puede hacer en $O(1)$ si calculamos estas posiciones de antemano durante el preprocesamiento). Entonces la respuesta a la consulta es la suma $T1[p..q-1]$ menos la suma $T2[p..q-1]$ .

¿Por qué? Consideremos el segmento $[p;q]$ en el tour de Euler. Contiene todas las aristas del camino que necesitamos de $i$ a $j$ pero también contiene un conjunto de aristas que yacen en otros caminos desde $i$ . Sin embargo hay una gran diferencia entre las aristas que necesitamos y el resto de las aristas: las aristas que necesitamos aparecerán solo una vez en la dirección de ida, y todas las demás aristas aparecen dos veces: una en la dirección de ida y una en la opuesta. Por tanto, la diferencia $T1[p..q-1] - T2[p..q-1]$ nos dará la respuesta correcta (el menos uno es necesario porque de lo contrario capturaríamos una arista extra que sale del vértice $j$ ). La consulta de suma en el Árbol de Segmentos se ejecuta en $O(\log N)$ .

Responder el primer tipo de consulta (pintar una arista) es aún más fácil: solo hay que actualizar $T1$ y $T2$ , a saber, realizar una sola actualización del elemento que corresponde a nuestra arista (encontrar la arista en la lista, otra vez, es posible en $O(1)$ , si se realiza esta búsqueda durante el preprocesamiento). Una sola modificación en el Árbol de Segmentos se realiza en $O(\log N)$ .

Implementación

Aquí está la implementación completa de la solución, incluyendo el cálculo del 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() {
    // lectura del 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 () {
            // consulta para pintar la arista x;
            // si start = true, entonces la arista se pinta; si no, se quita
            // la pintura
            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 {
            // consultar el número de aristas pintadas en el camino entre v1 y v2
            int l = lca(v1, v2);
            int result = query(l, v1) + query(l, v2);
            // result - la respuesta a la consulta
        }
    }
}