NC19995 [HAOI2015]树上操作

题目链接

题目

题目描述

有一棵点数为 N 的树，以点 1 为根，且树点有边权。

然后有 M 个 操作，分为三种：

操作 1 ：把某个节点 x 的点权增加 a 。

操作 2 ：把某个节点 x 为根的子树中所有点的点权都增加 a 。

操作 3 ：询问某个节点 x 到根的路径中所有点的点权和。

输入描述

第一行包含两个整数N, M。表示点数和操作数。
接下来一行N个整数，表示树中节点的初始权值。
接下来N-1行每行三个正整数 fr, to ，表示该树中存在一条边 (fr, to) 。
再接下来M行，每行分别表示一次操作。其中第一个数表示该操作的种类（ 1-3 ），之后接这个操作的参数（ x 或者 x a ） 。

输出描述

对于每个询问操作，输出该询问的答案。答案之间用换行隔开。

示例1

输入

5 5
1 2 3 4 5
1 2
1 4
2 3
2 5
3 3
1 2 1
3 5
2 1 2
3 3

输出

6
9
13

备注

对于 100% 的数据， \(N,M\le 100000\) ，且所有输入数据的绝对值都不会超过 10^6 。

题解

知识点：DFS序，线段树。

这题可以用树剖写，是板题。这里用dfs序写一下。

转换为dfs序后，每个子树都有一个开始标志和结束标志。我们查询根节点的开始标志到目标节点开始标志的这一个区间，如果不属于根节点到自己路径上的其他点，会同时遇到开始和结束标志。因此，我们可以给开始和结束标志赋予相同权值，但一正一负，那么区间和时，若同时遇到，则和为 \(0\) ，等价于没有算这个节点的值，而最后结果就只有路径和。

因此，我们使用dfs序，用线段树维护。其中节点属性 \(cnt\) ，表示一个区间的开始标志与结束标志的数量差，用以更新权值时计算。

时间复杂度 \(O((n+m)\log n)\)

空间复杂度 \(O(n)\)

代码

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

struct Graph {
    struct edge {
        int v, nxt;
    };
    int idx;
    vector<int> h;
    vector<edge> e;

    Graph(int n = 0, int m = 0) { init(n, m); }

    void init(int n, int m) {
        idx = 0;
        h.assign(n + 1, 0);
        e.assign(m + 1, {});
    }

    void add(int u, int v) {
        e[++idx] = { v,h[u] };
        h[u] = idx;
    }
};

struct T {
    int cnt;
    ll sum;
    static T e() { return { 0,0 }; }
    friend T operator+(const T &a, const T &b) { return { a.cnt + b.cnt,a.sum + b.sum }; }
};
struct F {
    ll add;
    static F e() { return { 0 }; }
    T operator()(const T &x) { return { x.cnt,x.sum + add * x.cnt }; }
    F operator()(const F &g) { return { g.add + add }; }
};

template<class T, class F>
class SegmentTreeLazy {
    int n;
    vector<T> node;
    vector<F> lazy;

    void push_down(int rt) {
        node[rt << 1] = lazy[rt](node[rt << 1]);
        lazy[rt << 1] = lazy[rt](lazy[rt << 1]);
        node[rt << 1 | 1] = lazy[rt](node[rt << 1 | 1]);
        lazy[rt << 1 | 1] = lazy[rt](lazy[rt << 1 | 1]);
        lazy[rt] = F::e();
    }

    void update(int rt, int l, int r, int x, int y, F f) {
        if (r < x || y < l) return;
        if (x <= l && r <= y) return node[rt] = f(node[rt]), lazy[rt] = f(lazy[rt]), void();
        push_down(rt);
        int mid = l + r >> 1;
        update(rt << 1, l, mid, x, y, f);
        update(rt << 1 | 1, mid + 1, r, x, y, f);
        node[rt] = node[rt << 1] + node[rt << 1 | 1];
    }

    T query(int rt, int l, int r, int x, int y) {
        if (r < x || y < l) return T::e();
        if (x <= l && r <= y) return node[rt];
        push_down(rt);
        int mid = l + r >> 1;
        return query(rt << 1, l, mid, x, y) + query(rt << 1 | 1, mid + 1, r, x, y);
    }

public:
    SegmentTreeLazy(int _n = 0) { init(_n); }
    SegmentTreeLazy(const vector<T> &src) { init(src); }

    void init(int _n) {
        n = _n;
        node.assign(n << 2, T::e());
        lazy.assign(n << 2, F::e());
    }
    void init(const vector<T> &src) {
        assert(src.size() >= 2);
        init(src.size() - 1);
        function<void(int, int, int)> build = [&](int rt, int l, int r) {
            if (l == r) return node[rt] = src[l], void();
            int mid = l + r >> 1;
            build(rt << 1, l, mid);
            build(rt << 1 | 1, mid + 1, r);
            node[rt] = node[rt << 1] + node[rt << 1 | 1];
        };
        build(1, 1, n);
    }

    void update(int x, int y, F f) { update(1, 1, n, x, y, f); }

    T query(int x, int y) { return query(1, 1, n, x, y); }
};

const int N = 100007;
Graph g;
int a[N];

int dfscnt;
int L[N], R[N];
void dfs(int u, int fa) {
    L[u] = ++dfscnt;
    for (int i = g.h[u];i;i = g.e[i].nxt) {
        int v = g.e[i].v;
        if (v == fa) continue;
        dfs(v, u);
    }
    R[u] = ++dfscnt;
}

int main() {
    std::ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    int n, m;
    cin >> n >> m;
    g.init(n, n << 1);
    for (int i = 1;i <= n;i++) cin >> a[i];
    for (int i = 1;i <= n - 1;i++) {
        int u, v;
        cin >> u >> v;
        g.add(u, v);
        g.add(v, u);
    }

    dfs(1, 0);
    vector<T> a_src(dfscnt + 1);
    for (int i = 1;i <= n;i++) {
        a_src[L[i]] = { 1,a[i] };
        a_src[R[i]] = { -1,-a[i] };
    }
    SegmentTreeLazy<T, F> sgt(a_src);

    while (m--) {
        int op, x;
        cin >> op >> x;
        if (op == 1) {
            int val;
            cin >> val;
            sgt.update(L[x], L[x], { val });
            sgt.update(R[x], R[x], { val });
        }
        else if (op == 2) {
            int val;
            cin >> val;
            sgt.update(L[x], R[x], { val });
        }
        else {
            cout << sgt.query(L[1], L[x]).sum << '\n';
        }
    }
    return 0;
}