[树链剖分] HDU 6962 I love tree

题目大意

有一棵 (n) 个点的数以及 (q) 个操作，每个点的点权初始时为0，操作1将从 (a) 到 (b) 的路径上点的权值分别加上 (1,4,9,16,dots)，操作2询问一个点的点权。((1leq n,qleq 10^5))

题解

首先肯定可以用树链剖分来维护。

注意到 (1,4,9,16,dots) 差分后是一个等差数列 (1,3,5,7,dots)，所以线段树只要实现区间加等差数列即可。在线段树的每个结点维护 (a,d) 两个标记，分别代表该区间要加上的等差数列的首项和公差，合并两个区间的这两个标记显然就是直接相加。设 (c=mathrm{lca}(a,b))，则 (asim c) 和 (csim b) 这两段加上的等差数列是不同的，需要分类讨论下，同时注意 (b) 点要加上真实值，在 (c) 的父亲处要消除影响。因为我们加的是差分后的数列，所以求一个点的权值只需要求以该点为根的子树内的所有数之和即可。时间复杂度 (O(nlog^2 n))。

Code

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

template<typename elemType>
inline void Read(elemType& T) {
    elemType X = 0, w = 0; char ch = 0;
    while (!isdigit(ch)) { w |= ch == '-';ch = getchar(); }
    while (isdigit(ch)) X = (X << 3) + (X << 1) + (ch ^ 48), ch = getchar();
    T = (w ? -X : X);
}

#define LL long long
#define int long long
const int maxn = 100010;
struct edge { int Next, to; };
edge G[maxn << 1];
int head[maxn], Deep[maxn], Fa[maxn], Hson[maxn], Size[maxn];
int Top[maxn], ID[maxn], TID[maxn];
LL Data[maxn];
int N, M, Root, Index = 0, cnt = 2;

inline void Add_Edge(int u, int v) {
    G[cnt].to = v;
    G[cnt].Next = head[u];
    head[u] = cnt++;
    return;
}

struct SegTreeNode {
    SegTreeNode() :Value(0LL), a1(0LL), d(0LL) {}
    LL Value, a1, d;
};
SegTreeNode SegTree[maxn << 2];

void DFSA(int u) {
    Size[u] = 1;
    int MaxSize = 0;
    for (int i = head[u];i;i = G[i].Next) {
        int v = G[i].to;
        if (v == Fa[u]) continue;
        Deep[v] = Deep[u] + 1;Fa[v] = u;
        DFSA(v);
        Size[u] += Size[v];
        if (Size[v] > MaxSize) { MaxSize = Size[v];Hson[u] = v; }
    }
    return;
}

void DFSB(int u, int Anc) {
    Top[u] = Anc;ID[u] = ++Index;TID[Index] = u;
    if (Hson[u]) DFSB(Hson[u], Anc);
    for (int i = head[u];i;i = G[i].Next) {
        int v = G[i].to;
        if (v == Fa[u] || v == Hson[u]) continue;
        DFSB(v, v);
    }
    return;
}

void Push_Up(int root) {
    SegTree[root].Value = SegTree[root << 1].Value + SegTree[root << 1 | 1].Value;
}

inline LL f(LL a1, LL d, LL n) { return n * (n - 1) / 2 * d + a1 * n; }

void Push_Down(int root, int L, int R) {
    if (!SegTree[root].d && !SegTree[root].a1) return;
    LL d = SegTree[root].d;
    LL a1 = SegTree[root].a1;
    SegTree[root].a1 = SegTree[root].d = 0;
    int mid = (L + R) >> 1;
    SegTree[root << 1].a1 += a1;
    SegTree[root << 1].d += d;
    SegTree[root << 1].Value += f(a1, d, mid - L + 1);
    SegTree[root << 1 | 1].a1 += a1 + d * (mid - L + 1);
    SegTree[root << 1 | 1].d += d;
    SegTree[root << 1 | 1].Value += f(a1 + d * (mid - L + 1), d, R - mid);
}

void Build_SegTree(int root, int L, int R) {
    if (L == R) { SegTree[root].Value = 0;return; }
    int mid = (L + R) >> 1;
    Build_SegTree(root << 1, L, mid);
    Build_SegTree(root << 1 | 1, mid + 1, R);
    Push_Up(root);
    return;
}

LL Update(int root, int L, int R, int UL, int UR, LL a1, LL d) {
    if (R < UL || UR < L) return 0;
    if (UL <= L && R <= UR) {
        SegTree[root].Value += f(a1, d, R - L + 1);
        SegTree[root].a1 += a1;
        SegTree[root].d += d;
        return R - L + 1;
    }
    Push_Down(root, L, R);
    int mid = (L + R) >> 1;
    LL e = Update(root << 1, L, mid, UL, UR, a1, d);
    e += Update(root << 1 | 1, mid + 1, R, UL, UR, a1 + d * e, d);
    Push_Up(root);
    return e;
}

LL Query(int root, int L, int R, int QL, int QR) {
    if (R < QL || QR < L) return 0LL;
    if (QL <= L && R <= QR) return SegTree[root].Value;
    Push_Down(root, L, R);
    int mid = (L + R) >> 1;
    return Query(root << 1, L, mid, QL, QR) +
        Query(root << 1 | 1, mid + 1, R, QL, QR);
}

int LCA(int u, int v) {
    while (Top[u] != Top[v]) {
        if (Deep[Top[u]] < Deep[Top[v]]) swap(u, v);
        u = Fa[Top[u]];
    }
    if (Deep[u] < Deep[v]) swap(u, v);
    return v;
}

void Maintain_Line(int u, int v) {
    int c = LCA(u, v);
    LL st = 1, num;
    int uu = u, vv = v;
    while (Top[u] != Top[c]) {
        st += (ID[u] - ID[Top[u]] + 1) * 2;
        Update(1, 1, N, ID[Top[u]], ID[u], st - 2, -2);
        u = Fa[Top[u]];
    }
    st += (ID[u] - ID[c] + 1) * 2;
    Update(1, 1, N, ID[c], ID[u], st - 2, -2);
    u = uu;
    Update(1, 1, N, ID[v], ID[v], (Deep[u] + Deep[v] - 2 * Deep[c] + 1) * (Deep[u] + Deep[v] - 2 * Deep[c] + 1), 0);
    Update(1, 1, N, ID[c], ID[c], -(Deep[u] - Deep[c] + 1) * (Deep[u] - Deep[c] + 1), 0);
    v = Fa[v];
    if (!v) return;
    if (Deep[v] < Deep[c]) {
        Update(1, 1, N, ID[v], ID[v], -((Deep[u] - Deep[c] + 1) * (Deep[u] - Deep[c] + 1)), 0);
        return;
    }
    st = 1 + (Deep[u] + Deep[v] - 2 * Deep[c] + 1) * 2;
    while (Top[v] != Top[c]) {
        st -= (ID[v] - ID[Top[v]] + 1) * 2;
        Update(1, 1, N, ID[Top[v]], ID[v], -(st + 2), -2);
        v = Fa[Top[v]];
    }
    st -= (ID[v] - ID[c] + 1) * 2;
    Update(1, 1, N, ID[c], ID[v], -(st + 2), -2);
    if (!Fa[c]) return;
    Update(1, 1, N, ID[Fa[c]], ID[Fa[c]], -((Deep[u] - Deep[c] + 1) * (Deep[u] - Deep[c] + 1)), 0);
}

signed main() {
    Root = 1;
    Read(N);
    if (N == 1) {
        Read(M);
        LL x = 0;
        while (M--) {
            int opt, u, v;
            Read(opt);
            if (opt == 1) { Read(u); Read(v); ++x; }
            else { Read(v); printf("%lld
", x); }
        }
        return 0;
    }
    for (int i = 1;i <= N - 1;++i) {
        int u, v;
        Read(u);Read(v);
        Add_Edge(u, v);
        Add_Edge(v, u);
    }
    DFSA(Root);DFSB(Root, Root);
    Build_SegTree(1, 1, N);
    Read(M);
    for (int i = 1;i <= M;++i) {
        int opt;
        scanf("%d", &opt);
        if (opt == 1) {
            int u, v;
            Read(u);Read(v);
            Maintain_Line(u, v);
        }
        else if (opt == 2) {
            int u;Read(u);
            printf("%lld
", Query(1, 1, N, ID[u], ID[u] + Size[u] - 1));
        }
    }

    return 0;
}