NC200179 Colorful Tree

题目链接

题目

题目描述

A tree structure with some colors associated with its vertices and a sequence of commands on it are given. A command is either an update operation or a query on the tree. Each of the update operations changes the color of a specified vertex, without changing the tree structure. Each of the queries asks the number of edges in the minimum connected subgraph of the tree that contains all the vertices of the specified color.
Your task is to find answers of each of the queries, assuming that the commands are performed in the given order.

输入描述

The input consists of a single test case of the following format.
n
\(a_1\) \(b_1\)
. . .
\(a_{n−1}\) \(b_{n−1}\)
\(c_1 ... c_n\)
m
\(command_1\)
. . .

\(command_m\)
The first line contains an integer n \((2 \leq n \leq 100000)\) , the number of vertices of the tree. The vertices are numbered 1 through n. Each of the following n−1 lines contains two integers \(a_i (1 \leq a_i \leq n)\) and \(b_i (1 \leq b_i \leq n)\) , meaning that the i-th edge connects vertices \(a_i\) and \(b_i\) . It is ensured that all the vertices are connected, that is, the given graph is a tree. The next line contains n integers, \(c_1\) through \(c_n\) , where \(c_j (1 \leq c_j \leq 100000)\) is the initial color of vertex j. The next line contains an integer m \((1 \leq m \leq 100000)\) , which indicates the number of commands. Each of the following m lines contains a command in the following format.
U \(x_k\) \(y_k\)
or
Q \(y_k\)
When the k-th command starts with U, it means an update operation changing the color of vertex \(x_k (1 \leq x_k \leq n)\) to \(y_k (1 \leq y_k \leq 100000)\) . When the k-th command starts with Q, it
means a query asking the number of edges in the minimum connected subgraph of the tree that contains all the vertices of color \(y_k (1 \leq y_k \leq 100000)\) .

输出描述

For each query, output the number of edges in the minimum connected subgraph of the tree containing all the vertices of the specified color. If the tree doesn’t contain any vertex of the specified color, output -1 instead.

示例1

输入

5 
1 2 
2 3 
3 4 
2 5 
1 2 1 2 3 
11 
Q 1 
Q 2 
Q 3 
Q 4 
U 5 1 
Q 1 
U 3 2 
Q 1 
Q 2 
U 5 4 
Q 1

输出

2 
2 
0 
-1 
3 
2 
2 
0

题解

知识点：DFS序，LCA，STL。

对于一种颜色，我们新加入一个点 \(x\) ，考虑其贡献。设dfn序左右最近的两个点 \(u,v\) ，那么其贡献为：

\[\frac{1}{2}(dist(u,x) + dist(v,x) - dist(u,v)) \]

可以分类讨论验证，首先一定是选择dfn序最近的两点，其次无论 \(x\) 位于 \(u,v\) 中间，还是一侧，这个式子都是满足的。其中对于一侧情况，我们统一取第一个点和最后一个点，可以同时满足左侧和右侧以及只有一个点的情况。

因此，我们对每个颜色的点建一个 set ，排序规则采用dfn序从小到大，更新查找时使用二分即可，增加删除可以写在同一个函数里。对于求距离，使用倍增LCA即可。

初始时，将每个点都加一遍。

时间复杂度 \(O((n+m)\log n)\)

空间复杂度 \(O(n \log n)\)

代码

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

struct Graph {
    struct edge {
        int v, nxt;
    };
    int idx;
    vector<int> h;
    vector<edge> e;

    Graph(int n = 0, int m = 0) { init(n, m); }

    void init(int n, int m) {
        idx = 0;
        h.assign(n + 1, 0);
        e.assign(m + 1, {});
    }

    void add(int u, int v) {
        e[++idx] = { v,h[u] };
        h[u] = idx;
    }
};

const int N = 100007;
Graph g;

namespace LCA {
    const int lgN = 16;
    int dep[N], f[lgN + 7][N];
    void dfs(int u, int fa = 0) {
        f[0][u] = fa;
        dep[u] = dep[fa] + 1;
        for (int i = 1;i <= lgN;i++) f[i][u] = f[i - 1][f[i - 1][u]];
        for (int i = g.h[u];i;i = g.e[i].nxt) {
            int v = g.e[i].v;
            if (v == fa) continue;
            dfs(v, u);
        }
    }

    int LCA(int u, int v) {
        if (dep[u] < dep[v]) swap(u, v);
        for (int i = lgN;i >= 0;i--) {
            if (dep[f[i][u]] >= dep[v]) u = f[i][u];
            if (u == v) return u;
        }
        for (int i = lgN;i >= 0;i--) {
            if (f[i][u] != f[i][v]) {
                u = f[i][u];
                v = f[i][v];
            }
        }
        return f[0][u];
    }

    int dis(int u, int v) { return dep[u] + dep[v] - 2 * dep[LCA(u, v)]; }
}

int dfncnt;
int rdfn[N];
void dfs(int u, int fa) {
    rdfn[u] = ++dfncnt;
    for (int i = g.h[u];i;i = g.e[i].nxt) {
        int v = g.e[i].v;
        if (v == fa) continue;
        dfs(v, u);
    }
}

struct cmp {
    bool operator()(int a, int b) const {
        return rdfn[a] < rdfn[b];
    }
};

int c[N];
set<int, cmp> st[N];
int ans[N];

void update(int x, int f) {
    if (f == -1) st[c[x]].erase(x);
    if (st[c[x]].size()) {
        auto it = st[c[x]].lower_bound(x);
        auto it1 = st[c[x]].begin();
        auto it2 = prev(st[c[x]].end());
        if (it != st[c[x]].begin() && it != st[c[x]].end()) {
            it1 = it;
            it2 = prev(it);
        }
        ans[c[x]] += f * (
            LCA::dis(*it1, x)
            + LCA::dis(*it2, x)
            - LCA::dis(*it1, *it2)
            ) / 2;
    }
    if (f == 1) st[c[x]].insert(x);
}

int main() {
    std::ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    int n;
    cin >> n;
    g.init(n, n << 1);
    for (int i = 1;i <= n - 1;i++) {
        int u, v;
        cin >> u >> v;
        g.add(u, v);
        g.add(v, u);
    }
    dfs(1, 0);
    LCA::dfs(1);

    for (int i = 1;i <= n;i++)
        cin >> c[i], update(i, 1);

    int m;
    cin >> m;
    while (m--) {
        char op;
        cin >> op;
        if (op == 'Q') {
            int col;
            cin >> col;
            cout << (st[col].size() ? ans[col] : -1) << '\n';
        }
        else {
            int x, col;
            cin >> x >> col;
            update(x, -1);
            c[x] = col;
            update(x, 1);
        }
    }
    return 0;
}