洛谷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;
}