트리의 간선 칠하기¶
이는 꽤 흔한 문제입니다. $N$개의 정점을 가진 트리 $G$가 주어집니다. 쿼리는 두 종류가 있습니다. 첫 번째는 간선 하나를 칠하는 것이고, 두 번째는 두 정점 사이에 있는 칠해진 간선의 개수를 묻는 것입니다.
여기서는 각 쿼리에 $O(\log N)$ 시간에 답하는 꽤 간단한 풀이(세그먼트 트리를 사용합니다)를 설명합니다. 전처리 단계는 $O(N)$ 시간이 걸립니다.
알고리즘¶
먼저, 두 번째 종류의 각 쿼리 $(i,j)$를 두 쿼리 $(l,i)$와 $(l,j)$로 바꾸기 위해 LCA를 구해야 합니다. 여기서 $l$은 $i$와 $j$의 LCA입니다. 쿼리 $(i,j)$의 답은 두 부분 쿼리의 답을 더한 값입니다. 이 두 쿼리는 모두 특별한 구조를 가지는데, 첫 번째 정점이 두 번째 정점의 조상이라는 것입니다. 이 글의 나머지 부분에서는 이러한 특별한 종류의 쿼리만 다룹니다.
먼저 전처리 단계부터 설명하겠습니다. 트리의 루트에서 깊이 우선 탐색을 실행하고, 이 깊이 우선 탐색의 오일러 투어를 기록합니다(각 정점은 탐색이 그 정점을 처음 방문할 때와 그 자식 중 하나에서 돌아올 때마다 리스트에 추가됩니다). 같은 기법을 LCA 전처리에서도 사용할 수 있습니다.
이 리스트에는 모든 간선이 들어 있으며(즉, $i$와 $j$가 간선의 양 끝점이라면 리스트 안에 $i$와 $j$가 서로 이웃한 위치가 존재한다는 의미입니다), 각 간선은 정확히 두 번 나타납니다. 한 번은 정방향($i$에서 $j$로, 여기서 정점 $i$가 정점 $j$보다 루트에 더 가깝습니다)으로, 한 번은 역방향($j$에서 $i$로)으로 나타납니다.
이 간선들에 대해 두 개의 리스트를 만듭니다. 첫 번째 리스트는 정방향인 모든 간선의 색을 저장하고, 두 번째 리스트는 역방향인 모든 간선의 색을 저장합니다. 간선이 칠해져 있으면 $1$을, 그렇지 않으면 $0$을 사용합니다. 이 두 리스트 위에 각각 세그먼트 트리(단일 원소 수정을 지원하는 합 세그먼트 트리)를 만들고, 이를 $T1$과 $T2$라고 부릅시다.
$i$가 $j$의 조상인 $(i,j)$ 형태의 쿼리에 답해 봅시다. $i$와 $j$ 사이의 경로에 칠해진 간선이 몇 개인지 구해야 합니다. 오일러 투어에서 $i$와 $j$가 처음 나타나는 위치를 찾고, 이를 $p$와 $q$라고 합시다(전처리 중에 이 위치들을 미리 계산해 두면 $O(1)$에 할 수 있습니다). 그러면 쿼리의 답은 합 $T1[p..q-1]$에서 합 $T2[p..q-1]$을 뺀 값입니다.
왜 그럴까요? 오일러 투어에서 구간 $[p;q]$를 생각해 봅시다. 이 구간은 $i$에서 $j$로 가는, 우리가 필요로 하는 경로의 모든 간선을 포함하지만, $i$에서 뻗어 나가는 다른 경로들 위에 있는 간선들도 포함합니다. 하지만 우리가 필요로 하는 간선과 나머지 간선 사이에는 한 가지 큰 차이가 있습니다. 우리가 필요로 하는 간선은 정방향으로 한 번만 나열되고, 다른 모든 간선은 정방향으로 한 번, 역방향으로 한 번, 총 두 번 나타납니다. 따라서 차 $T1[p..q-1] - T2[p..q-1]$이 올바른 답을 줍니다(1을 빼는 것이 필요한 이유는, 그렇지 않으면 정점 $j$에서 나가는 간선 하나를 추가로 포함하게 되기 때문입니다). 세그먼트 트리에서 합 쿼리는 $O(\log N)$에 수행됩니다.
첫 번째 종류의 쿼리(간선 칠하기)에 답하는 것은 더 쉽습니다. $T1$과 $T2$를 갱신하기만 하면 되는데, 즉 해당 간선에 대응하는 원소를 한 번 갱신하면 됩니다(리스트에서 간선을 찾는 것 역시 전처리 중에 이 탐색을 수행해 두면 $O(1)$에 가능합니다). 세그먼트 트리에서 단일 원소 수정은 $O(\log N)$에 수행됩니다.
구현¶
다음은 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() {
// reading the graph
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 () {
// request for painting edge x;
// if start = true, then the edge is painted, otherwise the painting
// is removed
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 {
// query the number of colored edges on the path between v1 and v2
int l = lca(v1, v2);
int result = query(l, v1) + query(l, v2);
// result - the answer to the request
}
}
}