P3833 [SHOI2012]魔法树（洛谷）【树链剖分+线段树模板】

P3833 [SHOI2012]魔法树（洛谷）【树链剖分+线段树模板】

思路

纯板子题，没有思维含量。但是不知道为什么数组要开2e5，开小了会报TLE，非常迷惑。

代码

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

typedef long long ll;

const int N = 2e5 + 10;

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

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

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

    void pushdown(int u)
    {
        auto &root = tr[u],
              &lson = tr[u << 1],
               &rson = tr[u << 1 | 1];
        if(!root.add)
            return ;
        lson.sum += (lson.r - lson.l + 1) * root.add;
        rson.sum += (rson.r - rson.l + 1) * root.add;
        lson.add += root.add;
        rson.add += root.add;
        root.add = 0;
    }

    void modify(int u, int l, int r, int d)
    {
        if(tr[u].l >= l && tr[u].r <= r)
        {
            tr[u].sum += (tr[u].r - tr[u].l + 1) * d;
            tr[u].add += d;
            return ;
        }
        pushdown(u);
        int 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(int u, int l, int r)
    {
        if(tr[u].l >= l && tr[u].r <= r)
            return tr[u].sum;
        pushdown(u);
        int mid = (tr[u].l + tr[u].r) >> 1;
        ll ans = 0;
        if(l <= mid)
            ans = query(u << 1, l, r);
        if(r > mid)
            ans += query(u << 1 | 1, l, r);
        return ans;
    }
};

struct tree_chain
{
    line_tree line;
    int head[N], nxt[N], es[N * 2];
    int wson[N], sz[N], dfsn[N], par[N], dep[N], top[N];
    int idx, tot, root;

    void init(int n)
    {
        line.build(1, 1, n);
        root = 0;
        memset(head, -1, sizeof head);
    }

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

    void dfs1(int u, int fa)
    {
        sz[u] = 1;
        for(int i = head[u]; ~i; i = nxt[i])
        {
            int v = es[i];
            if(v == fa)
                continue;
            dep[v] = dep[u] + 1;
            par[v] = u;
            dfs1(v, u);
            sz[u] += sz[v];
            if(sz[wson[u]] < sz[v])
                wson[u] = v;
        }
    }

    void dfs2(int u, int tp)
    {
        dfsn[u] = ++tot;
        top[u] = tp;
        if(wson[u])
            dfs2(wson[u], tp);
        for(int i = head[u]; ~i; i = nxt[i])
        {
            int v = es[i];
            if(v == par[u] || v == wson[u])
                continue;
            dfs2(v, v);
        }
    }

    void modify(int a, int b, int d)
    {
        while(top[a] != top[b])
        {
            if(dep[top[a]] < dep[top[b]])
                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);
    }

    ll query(int u)
    {
        return line.query(1, dfsn[u], dfsn[u] + sz[u] - 1);
    }
}chain;

int main()
{
    int n;
    scanf("%d", &n);
    chain.init(n);
    for(int i = 1, u, v; i < n; i++)
    {
        cin >> u >> v;
        chain.addEdge(u + 1, v + 1);
        chain.addEdge(v + 1, u + 1);
    }
    chain.dfs1(1, -1);
    chain.dfs2(1, 1);
    int q;
    scanf("%d", &q);
    char op;
    for(int i = 0, u, v, d; i < q; i++)
    {
        scanf("\n%c", &op);
        if(op == 'A')
        {
            scanf("%d%d%d", &u, &v, &d);
            chain.modify(u + 1, v + 1, d);
        }
        else
        {
            scanf("%d", &u);
            printf("%lld\n", chain.query(u + 1));
        }
    }
    return 0;
}