14243: 【原4243】DFS序列

题目

题目描述

author: 青铜角斗士原OJ链接：https://acm.sjtu.edu.cn/OnlineJudge-old/problem/4243

Description

给一棵有根树，这棵树由编号为1 ... N的N个结点组成。根结点的编号为R。每个结点都有一个权值，结点i的权值为$v_i$。接下来有M组操作，操作分为两类：

1 a x，表示将结点a的子树上所有结点的权值增加；

2 a，表示求结点a的子树上所有结点的权值之和。

Input Format

第一行有三个整数N, M和R。

第二行有N个整数，第i个整数表示$v_i$。

在接下来的N-1行中，每行两个整数，表示一条边。

在接下来的M行中，每行一组操作。

Output Format

对于每组2 a操作，输出一个整数，表示「以结点a为根的子树」上所有结点的权值之和。

Sample Input

10 14 9
12 -6 -4 -3 12 8 9 6 6 2
8 2
2 10
8 6
2 7
7 1
6 3
10 9
2 4
10 5
1 4 -1
2 2
1 7 -1
2 10
1 10 5
2 1
1 7 -5
2 5
1 1 8
2 7
1 8 8
2 2
1 5 5
2 6

Sample Output

Subtasks

对于60%的数据，满足$1\leq N,M \leq 10^3, 1 \leq R \leq N, -10^6 \leq v_i, x \leq 10^6$。
对于100%的数据，满足$1\leq N,M \leq 10^6, 1 \leq R \leq N, -10^6 \leq v_i, x \leq 10^6$。

ligongzzz's solution

#include "iostream"
#include "vector"
#include "stack"
using namespace std;

class node {
public:
    long long val = 0;
    long long parent = 0;
    long long num = 0;
    long long tag = 0;
    vector<long long> edges;
};

long long n, m, r;

long long find_parent(long long pos,vector<node>& node_list) {
    long long ans = 1;
    for (auto p : node_list[pos].edges) {
        if (p != node_list[pos].parent) {
            node_list[p].parent = pos;
            ans += find_parent(p, node_list);
            node_list[pos].val += node_list[p].val;
        }
    }
    node_list[pos].num = ans;
    return ans;
}

void push_down(long long pos, vector<node>& node_list) {
    stack<long long> stack_data;
    for (auto p = node_list[pos].parent; p; p = node_list[p].parent) {
        stack_data.push(p);
    }
    while (!stack_data.empty()) {
        auto temp = stack_data.top();
        stack_data.pop();
        if (node_list[temp].tag) {
            for (auto p : node_list[temp].edges) {
                if (p != node_list[temp].parent) {
                    node_list[p].val += node_list[temp].tag * node_list[p].num;
                    node_list[p].tag += node_list[temp].tag;
                }
            }
            node_list[temp].tag = 0;
        }
    }
}

void push_up(long long pos, long long val, vector<node>& node_list) {
    for (pos = node_list[pos].parent; pos; pos = node_list[pos].parent) {
        node_list[pos].val += val;
    }
}

void update(long long pos, long long val,vector<node>& node_list) {
    push_down(pos, node_list);
    node_list[pos].tag += val;
    node_list[pos].val += node_list[pos].num * val;
    push_up(pos, node_list[pos].num * val, node_list);
}

long long query(long long pos, vector<node>& node_list) {
    push_down(pos, node_list);
    return node_list[pos].val;
}

int main() {
    scanf("%lld %lld %lld", &n, &m, &r);

    vector<node> node_list(n + 1);
    for (long long i = 1; i <= n; ++i) {
        scanf("%lld", &node_list[i].val);
    }

    for (long long i = 1; i <= n - 1; ++i) {
        long long a, b;
        scanf("%lld %lld", &a, &b);
        node_list[a].edges.push_back(b);
        node_list[b].edges.push_back(a);
    }

    find_parent(r, node_list);

    for (; m > 0; --m) {
        long long op;
        scanf("%lld", &op);
        if (op == 1) {
            long long a, x;
            scanf("%lld %lld", &a, &x);
            update(a, x, node_list);
        }
        else {
            long long a;
            scanf("%lld", &a);
            cout << query(a, node_list) << "\n";
        }
    }

    return 0;
}