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)$.
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)$.
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
}
}
}