P3178 [HAOI2015]树上操作（洛谷）【树链剖分+线段树模板】

P3178 [HAOI2015]树上操作（洛谷）【树链剖分+线段树模板】

思路

这题是一个树链剖分模板题，本来是参照秦淮岸灯火阑珊的博客树链剖分(轻重链剖分)算法笔记自己写一遍树链剖分并作为以后的板子，结果非常不幸纠结了一整天才搞完。原博客的代码有点乱，我一直以为自己错是因为有细节错误，结果找了很久发现原博客的代码有 #define int long long ……于是把所有的int全部换long long就过了……

8.11更新：modify那里有两步写错了，居然还能过……我在代码中标注了，如果是想用这题测板子的同学，不要来，可以去做魔法树这题。

代码

#include <bits/stdc++.h>
#define pb push_back
using namespace std;
typedef long long ll;
const ll INF = 0x3f3f3f3f;
const ll N = 1e6 + 10;

ll es[N * 2], head[N], nxt[N * 2], w[N * 2], idx;
ll dep[N], sz[N], top[N], dfsn[N], par[N], wson[N], old[N], tot;

struct line_tree
{
    struct node
    {
        ll l, r;
        ll sum, add;
    } tr[N * 4];

    void pushup(ll u)
    {
        tr[u].sum = tr[u << 1].sum + tr[u << 1 | 1].sum;
    }

    void build(ll u, ll l, ll r)
    {
        if(l == r)
        {
            tr[u] = {l, r, w[old[l]], 0};
            return ;
        }
        tr[u] = {l, r, 0, 0};
        ll mid = (l + r) >> 1;
        build(u << 1, l, mid);
        build(u << 1 | 1, mid + 1, r);
        pushup(u);
    }

    void pushdown(ll u)
    {
        ll add = tr[u].add;
        tr[u << 1].sum += (tr[u << 1].r - tr[u << 1].l + 1) * add;
        tr[u << 1].add += add;
        tr[u << 1 | 1].sum += (tr[u << 1 | 1].r - tr[u << 1 | 1].l + 1) * add;
        tr[u << 1 | 1].add += add;
        tr[u].add = 0;
    }

    void modify(ll u, ll l, ll r, ll d)
    {
        if(tr[u].l >= l && tr[u].r <= r)
        {
            tr[u].add += d;
            tr[u].sum += (tr[u].r - tr[u].l + 1) * d;
            return ;
        }
        pushdown(u);
        ll mid = (tr[u].l + tr[u].r) >> 1;
        if(l <= mid)
            modify(u << 1, l, r, d);
        if(r > mid)
            modify(u << 1 | 1, l, r, d);
        pushup(u);
    }

    ll query(ll u, ll l, ll r)
    {
        if(tr[u].l >= l && tr[u].r <= r)
            return tr[u].sum;
        pushdown(u);
        ll mid = (tr[u].r + tr[u].l) >> 1;
        ll res = 0;
        if(l <= mid)
            res = query(u << 1, l, r);
        if(r > mid)
            res += query(u << 1 | 1, l, r);
        return res;
    }
} line;

void init()
{
    memset(head, -1, sizeof head);
}

void addEdge(ll u, ll v)
{
    es[++idx] = v;
    nxt[idx] = head[u];
    head[u] = idx;
}

void dfs1(ll u, ll fa) // 当前节点与父亲节点
{
    sz[u] = 1;
    for(ll i = head[u]; ~i; i = nxt[i])
    {
        ll v = es[i];
        if(v == fa)
            continue;
        dep[v] = dep[u] + 1;
        par[v] = u;
        dfs1(v, u);
        sz[u] += sz[v];
        if(sz[v] > sz[wson[u]])
            wson[u] = v;
    }
}

void dfs2(ll u, ll tp) //当前节点与链头节点
{
    dfsn[u] = ++tot;
    old[tot] = u;
    top[u] = tp;
    if(wson[u])
        dfs2(wson[u], tp);
    for(ll i = head[u]; ~i; i = nxt[i])
    {
        ll v = es[i];
        if(v == wson[u] || v == par[u])
            continue;
        dfs2(v, v);
    }
}

void modify(ll x, ll v) //给子树x全部加上v
{
    line.modify(1, dfsn[x], dfsn[x] + sz[x] - 1, v);
}

ll query(ll a, ll b)
{
    ll ans = 0;
    while(top[a] != top[b])
    {
        if(dep[top[a]] < dep[top[b]])
            swap(a, b);
        ans += line.query(1, dfsn[top[a]], dfsn[a]);
        a = par[top[a]];
    }
    if(dep[a] < dep[b])
        swap(a, b);
    ans += line.query(1, dfsn[top[a]], dfsn[a]);
    return ans;
}

void modify(ll a, ll b, ll d)
{
    while(top[a] != top[b])
    {
        if(dep[top[a]] < dep[top[b]]) //这里小心，必须找top
            swap(a, b);
        line.modify(1, dfsn[top[a]], dfsn[a], d);
        a = par[top[a]];
    }
    if(dep[a] < dep[b])
        swap(a, b);
    line.modify(1, dfsn[a], dfsn[b], d); //这里注意不要照抄上面的modify
}

int main()
{
    ll n, m;
    scanf("%lld%lld", &n, &m);
    init();
    for(ll i = 1; i <= n; i++)
    {
        scanf("%lld", &w[i]);
    }
    for(ll i = 0, u, v; i < n - 1; i++)
    {
        scanf("%lld%lld", &u, &v);
        addEdge(u, v);
        addEdge(v, u);
    }
    dfs1(1, -1);
    dfs2(1, 1);
    line.build(1, 1, n);
    while(m--)
    {
        ll op, x, a;
        scanf("%lld%lld", &op, &x);
        if(op == 1)
        {
            scanf("%lld", &a);
            line.modify(1, dfsn[x], dfsn[x], a);
        }
        else if(op == 2)
        {
            scanf("%lld", &a);
            modify(x, a);
        }
        else
        {
            printf("%lld\n", query(1, x));
        }
    }
    return 0;
}