洛谷P7735 [NOI2021] 轻重边题解

题目链接：P7735 [NOI2021] 轻重边

题意：

小 W 有一棵 \(n\) 个结点的树，树上的每一条边可能是轻边或者重边。

在所有操作开始前，树上所有边都是轻边。

接下来你需要对树进行 \(m\) 次操作，操作有以下两种：

给定两个点 \(a\) 和 \(b\)，首先对于 \(a\) 到 \(b\) 路径上的所有点 \(x\)（包含 \(a\) 和 \(b\)），你要将与 \(x\) 相连的所有边变为轻边。然后再将 \(a\) 到 \(b\) 路径上包含的所有边变为重边。

给定两个点 \(a\) 和 \(b\)，你需要计算当前 \(a\) 到 \(b\) 的路径上一共包含多少条重边。

输入格式：

本题有多组数据，输入数据第一行一个正整数 \(T\)，表示数据组数。对于每组数据：

第一行包含两个整数 \(n\) 和 \(m\)，其中 \(n\) 表示结点数量，\(m\) 表示操作数量。

接下来 \(n - 1\) 行，每行包含两个整数 \(u,v\)，表示树上的一条边。

接下来 \(m\) 行，每行包含三个整数 \({\mathrm{op}}_i\ a_i\ b_i\)，描述一个操作，其中 \({\mathrm{op}}_i = 1\) 表示第 \(1\) 类操作，\({\mathrm{op}}_i = 2\) 表示第 \(2\) 类操作。

数据保证 \(a_i \neq b_i\)。

输出格式：

对于每一次第 \(2\) 类操作，输出一行一个整数表示答案。

数据范围：

对于所有测试数据：\(1\le T \le 3,~1 \le n, m \le {10}^5\)。

考虑将边的限制转化为点的限制。给每个点 \(i\) 设定初始权值为 \(i\) 。

那么每次修改就变成了，将路径 \(u \leadsto v\) 上的点修改为 \(o\) （其中 \(o\) 是时间戳，从 \(n+1\) 开始）。

显然，这样的好处是，一条边是重边，当且仅当这条边两端的点的权值相等。

由于是树上的路径修改，所以考虑树链剖分。

线段树每个节点维护答案和左右端点的权值。这样每次合并特判左儿子的右端点和右儿子的左端点即可。

时间复杂度 \(\mathcal{O}(n \log^2 n)\)

代码：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
void up(int &x, int y) { x < y ? x = y : 0; }
void down(int &x, int y) { x > y ? x = y : 0; }
#define rep(i, a, b) for(int i = (a), i##END = (b); i <= i##END; i++)
#define Rep(i, a, b) for(int i = (a), i##END = (b); i >= i##END; i--)
#define N ((int)(1e5 + 15))
#define ls(at) ((at) << 1)
#define rs(at) ((at) << 1 | 1)

vector<int> vec[N];
int n, o, tot, w[N], val[N];
int fa[N], son[N], sz[N], dep[N], top[N], dfn[N];
struct node { int l, r, ans, tag; } tr[N * 4];
void dfs1(int u, int f)
{
    fa[u] = f; dep[u] = dep[fa[u]] + 1; sz[u] = 1; son[u] = 0;
    for(int v : vec[u]) if(v != f)
    {
        dfs1(v, u); sz[u] += sz[v];
        if(sz[v] > sz[son[u]]) son[u] = v;
    }
}
void dfs2(int u, int topf)
{
    dfn[u] = ++tot; top[u] = topf; val[dfn[u]] = w[u];
    if(son[u]) dfs2(son[u], topf);
    for(int v : vec[u]) if(v != fa[u] && v != son[u]) dfs2(v, v);
}
void push_up(int at)
{
    tr[at].l = tr[ls(at)].l; tr[at].r = tr[rs(at)].r;
    tr[at].ans = tr[ls(at)].ans + tr[rs(at)].ans + (tr[rs(at)].l == tr[ls(at)].r);
}
void build(int l = 1, int r = n, int at = 1)
{
    if(l == r) { return tr[at] = {val[l], val[l], 0, 0}, void(0); }
    tr[at].tag = 0; int mid = (l + r) >> 1;
    build(l, mid, ls(at)); build(mid + 1, r, rs(at)); push_up(at);
}
void push_down(int l, int r, int at)
{
    if(!tr[at].tag) return;
    int mid = (l + r) >> 1;
    tr[ls(at)].ans = mid - l; tr[rs(at)].ans = r - mid - 1;
    tr[ls(at)].l = tr[ls(at)].r = tr[rs(at)].l = tr[rs(at)].r = tr[at].tag;
    tr[ls(at)].tag = tr[rs(at)].tag = tr[at].tag; tr[at].tag = 0;
}
void update(int nl, int nr, int k, int l = 1, int r = n, int at = 1)
{
    if(nl <= l && r <= nr)
    {
        tr[at].l = tr[at].r = tr[at].tag = k; tr[at].ans = r - l;
        return;
    }
    push_down(l, r, at); int mid = (l + r) >> 1;
    if(nl <= mid) update(nl, nr, k, l, mid, ls(at));
    if(nr > mid) update(nl, nr, k, mid + 1, r, rs(at));
    push_up(at);
}
int query(int nl, int nr, int l = 1, int r = n, int at = 1)
{
    if(nl <= l && r <= nr) return tr[at].ans;
    push_down(l, r, at); int mid = (l + r) >> 1, res = 0;
    if(nl <= mid) res += query(nl, nr, l, mid, ls(at));
    if(nr > mid) res += query(nl, nr, mid + 1, r, rs(at));
    if(nl <= mid && nr > mid && tr[ls(at)].r == tr[rs(at)].l) ++res;
    return res;
}
int query2(int x, int l = 1, int r = n, int at = 1)
{
    if(l == r) return tr[at].l;
    push_down(l, r, at); int mid = (l + r) >> 1;
    if(x <= mid) return query2(x, l, mid, ls(at));
    else return query2(x, mid + 1, r, rs(at));
}
void upRange(int x, int y)
{
    for(++o; top[x] != top[y]; x = fa[top[x]])
    {
        if(dep[top[x]] < dep[top[y]]) swap(x, y);
        update(dfn[top[x]], dfn[x], o);
    }
    if(dep[x] > dep[y]) swap(x, y);
    update(dfn[x], dfn[y], o);
}
int qRange(int x, int y)
{
    int res = 0;
    for(; top[x] != top[y]; x = fa[top[x]])
    {
        if(dep[top[x]] < dep[top[y]]) swap(x, y);
        res += query(dfn[top[x]], dfn[x]);
        if(query2(dfn[fa[top[x]]]) == query2(dfn[top[x]])) ++res;
    }
    if(dep[x] > dep[y]) swap(x, y);
    res += query(dfn[x], dfn[y]); return res;
}
void work()
{
    o = tot = 0; int q; cin >> n >> q;
    rep(i, 1, n) { w[i] = ++o; vec[i].clear(); }
    rep(i, 1, n - 1)
    {
        static int u, v; cin >> u >> v;
        vec[u].push_back(v); vec[v].push_back(u);
    }
    dfs1(1, 0); dfs2(1, 1); build();
    while(q--)
    {
        static int op, x, y; cin >> op >> x >> y;
        if(op == 1) upRange(x, y);
        else cout << qRange(x, y) << '\n';
    }
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    int qwq; cin >> qwq; while(qwq--) work();
    return 0;
}