题解 [ABC133F] Colorful Tree

· · 题解

题解怎么全是 poly log 啊。提供一种好想的根号做法。（虽然其实 poly log 貌似也很一眼

看到颜色数我们考虑根号分治。

对于出现次数很少的颜色，我们可以在询问的时候，先求出不改变边权时候两个点的距离，然后枚举每条边判断是否出现在询问的两点的路径上，如果出现就更改答案。

对于出现次数很多的颜色，那么颜色种类数很少，考虑对每种颜色预处理答案。我们发现我们的询问本质上是要求两个点之间将指定颜色边权改为 0 后的距离，以及求两个点之间指定颜色的个数。对于我们给定的颜色，这两个东西都可以通过一遍 DFS O(n) 预处理，然后询问的时候只需要求个 LCA 即可。

设出现次数 \ge B 的颜色为出现次数多的颜色，那么总复杂度 O(qB \log n+\dfrac{n^2}{B})。设 n,q 同阶，那么 B 取 \sqrt{\dfrac{n}{\log n}} 时复杂度达到 O(n \sqrt{n \log n})。如果使用 O(n \log n) - O(1) 的 LCA 那么块长取 \sqrt{n} 时复杂度为 O(n \sqrt n)。

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int B = 350;
struct Edge {
    int to, next;
    int c; ll w;
} e[200005];
int head[100005], len;
void Insert(int u, int v, int c, ll w) {
    if (!len) len = 1;
    e[++len] = { v, head[u], c, w };
    head[u] = len;
}
int n, q;
int prt[100005][20], dep[100005];
void PreLCA(int u, int fa) {
    dep[u] = dep[fa] + 1; prt[u][0] = fa;
    for (int i = 1; i < 20; i++) prt[u][i] = prt[prt[u][i - 1]][i - 1];
    for (int i = head[u]; i; i = e[i].next) {
        int v = e[i].to;
        if (v == fa) continue;
        PreLCA(v, u);
    }
}
int LCA(int u, int v) {
    if (dep[u] > dep[v]) swap(u, v);
    int d = dep[v] - dep[u];
    for (int i = 0; i < 20; i++) if (d >> i & 1) v = prt[v][i];
    if (u == v) return u;
    for (int i = 19; ~i; i--) {
        if (prt[u][i] != prt[v][i]) {
            u = prt[u][i]; v = prt[v][i];
        }
    }
    return prt[u][0];
}
int id[100005], iCnt;
ll dN[305][100005], dC[305][100005];
void DFS(int u, int fa, int col) {
    for (int i = head[u]; i; i = e[i].next) {
        int v = e[i].to;
        if (v == fa) continue;
        dC[id[col]][v] = dC[id[col]][u]; dN[id[col]][v] = dN[id[col]][u];
        if (e[i].c == col) dC[id[col]][v]++;
        else dN[id[col]][v] += e[i].w;
        DFS(v, u, col);
    }
}
int cCnt[100005];
vector<int> eC[100005];
bool PointOnPath(int u, int v, int x) { return (LCA(u, x) == x || LCA(v, x) == x) && dep[x] >= dep[LCA(u, v)]; }
bool EdgeOnPath(int u, int v, int i) { return PointOnPath(u, v, e[i].to) && PointOnPath(u, v, e[i ^ 1].to); }
int main() {
    scanf("%d%d", &n, &q);
    for (int i = 1; i < n; i++) {
        int u, v, c; ll w;
        scanf("%d%d%d%lld", &u, &v, &c, &w);
        Insert(u, v, c, w); Insert(v, u, c, w);
        cCnt[c]++; eC[c].emplace_back(len);
    }
    PreLCA(1, 0); DFS(1, 0, 0);
    for (int i = 1; i < n; i++) if (cCnt[i] >= B) id[i] = ++iCnt, DFS(1, 0, i);
    while (q--) {
        int x, u, v; ll y;
        scanf("%d%lld%d%d", &x, &y, &u, &v);
        if (id[x]) {
            ll diN = dN[id[x]][u] + dN[id[x]][v] - 2 * dN[id[x]][LCA(u, v)];
            ll diC = dC[id[x]][u] + dC[id[x]][v] - 2 * dC[id[x]][LCA(u, v)];
            printf("%lld\n", diN + diC * y);
        }
        else {
            ll ans = dN[0][u] + dN[0][v] - 2 * dN[0][LCA(u, v)];
            for (int i : eC[x]) {
                if (EdgeOnPath(u, v, i)) {
                    ans += y - e[i].w;
                }
            }
            printf("%lld\n", ans);
        }
    }
    return 0;
}