LeetCode LCP 05. 發 LeetCoin DFS序+帶懶惰標記的線段樹

題目描述

力扣決定給一個刷題團隊發 LeetCoin 作為獎勵。同時，為了監控給大家發了多少 LeetCoin，力扣有時候也會進行查詢。

該刷題團隊的管理模式可以用一棵樹表示：

1. 團隊只有一個負責人，編號為 1。除了該負責人外，每個人有且僅有一個領導（負責人沒有領導）；
2. 不存在循環管理的情況，如 A 管理 B，B 管理 C，C 管理 A。

力扣想進行的操作有以下三種：

1. 給團隊的一個成員（也可以是負責人）發一定數量的 LeetCoin；
2. 給團隊的一個成員（也可以是負責人），以及他/她管理的所有人（即他/她的下屬、他/她下屬的下屬，……），發一定數量的 LeetCoin；
3. 查詢某一個成員（也可以是負責人），以及他/她管理的所有人被發到的 LeetCoin 之和。

輸入

1. N 表示團隊成員的個數（編號為 1 ~ N，負責人為 1）；
2. leadership 是大小為 (N - 1) * 2 的二維數組，其中每個元素 [a, b] 代表 b 是 a 的下屬；
3. operations是一個長度為 Q 的二維數組，代表以時間排序的操作，格式如下：
   • operations[i][0] = 1: 代表第一種操作，operations[i][1]代表成員的編號，operations[i][2] 代表 LeetCoin 的數量；
   • operations[i][0] = 2: 代表第二種操作，operations[i][1] 代表成員的編號，operations[i][2] 代表 LeetCoin 的數量；
   • operations[i][0] = 3: 代表第三種操作，operations[i][1] 代表成員的編號；

輸出

返回一個數組，數組里是每次查詢的返回值（發 LeetCoin 的操作不需要任何返回值）。

由於發的 LeetCoin 很多，請把每次查詢的結果模 1e9+7 (1000000007)。

樣例

輸入：N = 6, leadership = [[1, 2], [1, 6], [2, 3], [2, 5], [1, 4]],
operations = [[1, 1, 500], [2, 2, 50], [3, 1], [2, 6, 15], [3, 1]]
輸出：[650, 665]
解釋：團隊的管理關系見下圖。
第一次查詢時，每個成員得到的LeetCoin的數量分別為（按編號順序）：500, 50, 50, 0, 50, 0;
第二次查詢時，每個成員得到的LeetCoin的數量分別為（按編號順序）：500, 50, 50, 0, 50, 15.

限制

1. 1 <= N <= 50000
2. 1 <= Q <= 50000
3. operations[i][0] != 3 時，1 <= operations[i][2] <= 5000

---

算法

(DFS 序 + 線段樹) \(O(n + Qlog n)\)

1. 首先，將樹結構轉化為一維序列結構，將子樹更新和子樹查詢轉為區間更新和區間查詢， DFS 序可以解決此問題。具體做法：維護一個時間戳，時間戳從 0 開始。當新遞歸進入一個結點時，時間戳加 1，且這個結點的 left 記錄為更新后的時間戳。然后深度優先遞歸遍歷子樹，當這個結點要回溯的時候，記錄該結點的 right 為當前的時間戳。這樣之后，每個結點就有了一個 left 值和一個 right 值。

   void dfs(int u, int& ts, const vector<vector<int>>& graph,
            vector<int>& left, vector<int> &right) {
       left[u] = ++ts;
       for (int v : graph[x])
           dfs(v, ts, graph, left, right);
   
       right[u] = ts;
   }
   

2. DFS 序以某個結點為根的子樹（包括自己）可以用子樹根結點的 left 和 right 值來表示。以樣例為例，整個子樹的 DFS 序為 (1, 2, 3, 5, 6, 4)，結點 2 的 left 為 2，right 為 4，所以區間 [2, 4] 就是以結點 2 為根的子樹。

3. 帶懶惰標記的線段樹可以實現區間更新與區間查詢。

時間復雜度

• 建立 DFS 序的時間復雜度為 \(O(n)\)。
• 線段樹的建立需要 \(O(n)\) 的時間復雜度，線段樹的單次更新和查詢需要 \(O(log n)\) 的時間復雜度。
• 故總時間復雜度為 \(O(n + Qlog n)\)。

空間復雜度

• 線段樹和其他數組占用的空間為 \(O(n)\)。

C++ 代碼

const int MOD = 1000000007;
const int N = 50010;
class Solution {
public:
    struct Node{
        int l, r;
        int sum, add;
    }tr[N << 2];

    void dfs(int x, int& ts, const vector<vector<int>>& graph, vector<int>& left, vector<int> &right) {
        left[x] = ++ts;
        for (int v : graph[x])
            dfs(v, ts, graph, left, right);

        right[x] = ts;
    }
    void build(int u, int l, int r){
        if(l == r) tr[u] = {l ,r, 0, 0};
        else {
            tr[u] = {l, r, 0, 0};
            int mid = (l + r) >> 1;
            build(u<<1, l, mid), build(u<<1|1, mid + 1, r);
            pushup(u);
        }
    }
    void pushdown(int u){
        if(!tr[u].add) return ;
        tr[u<<1].add = (tr[u<<1].add + tr[u].add) % MOD;
        tr[u<<1].sum =  (tr[u<<1].sum + (tr[u<<1].r - tr[u<<1].l + 1) * tr[u].add % MOD) % MOD;
        tr[u<<1|1].add = (tr[u<<1|1].add + tr[u].add) % MOD;
        tr[u<<1|1].sum =  (tr[u<<1|1].sum + (tr[u<<1|1].r - tr[u<<1|1].l + 1) * tr[u].add % MOD) % MOD;
        tr[u].add = 0;
        
    }
    void pushup(int u){
        tr[u].sum = (tr[u<<1].sum + tr[u<<1|1].sum) % MOD;
    }

    int 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;
        int sum = 0;
        if(l <= mid) sum = query(u<<1, l, r);
        if(r > mid) sum = (sum + query(u<<1|1, l, r)) % MOD;
        return sum;
    }

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

    vector<int> bonus(int n, vector<vector<int>>& leadership, vector<vector<int>>& operations) {
        vector<vector<int>> graph(n + 1);

        for (const auto &v : leadership)
            graph[v[0]].push_back(v[1]);

        vector<int> left(n + 1), right(n + 1);

        int ts = 0;
        dfs(1, ts, graph, left, right);
        
        build(1, 1, n);
        vector<int> ans;
        for (const auto &op : operations) {
            if (op[0] == 1) modify(1, left[op[1]], left[op[1]], op[2]);
            else if (op[0] == 2) modify(1, left[op[1]], right[op[1]], op[2]);
            else ans.push_back(query(1, left[op[1]], right[op[1]]));
        }
        return ans;
    }
};

指針寫法：

#define MOD 1000000007
#define LL long long

struct Node {
    int sum, lazy;
    Node *l, *r;
    Node(): sum(0), lazy(0), l(NULL), r(NULL){}
};

class Solution {
public:
    void dfs(int x, int& ts, const vector<vector<int>>& graph,
             vector<int>& left, vector<int> &right) {
        left[x] = ++ts;
        for (int v : graph[x])
            dfs(v, ts, graph, left, right);

        right[x] = ts;
    }

    Node* build(int L, int R) {
        Node *ret = new Node();
        if (L == R)
            return ret;

        int mid = (L + R) >> 1;
        ret -> l = build(L, mid);
        ret -> r = build(mid + 1, R);

        return ret;
    }

    void pushdown(Node *cur, int sz) {
        cur -> l -> lazy = (cur -> l -> lazy + cur -> lazy) % MOD;
        cur -> r -> lazy = (cur -> r -> lazy + cur -> lazy) % MOD;
        cur -> l -> sum =
            (cur -> l -> sum + (LL)(cur -> lazy) * (sz - (sz >> 1))) % MOD;
        cur -> r -> sum =
            (cur -> r -> sum + (LL)(cur -> lazy) * (sz >> 1)) % MOD;
        cur -> lazy = 0;
    }

    int query(int L, int R, int l, int r, Node *cur) {
        if (L <= l && r <= R)
            return cur -> sum;

        pushdown(cur, r - l + 1);

        int mid = (l + r) >> 1, ret = 0;
        if (L <= mid) ret = (ret + query(L, R, l, mid, cur -> l)) % MOD;
        if (mid < R) ret = (ret + query(L, R, mid + 1, r, cur -> r)) % MOD;

        return ret;
    }

    void update(int p, int x, int l, int r, Node *cur) {
        if (l == r) {
            cur -> sum = (cur -> sum + x) % MOD;
            return;
        }

        pushdown(cur, r - l + 1);
        int mid = (l + r) >> 1;
        if (p <= mid) update(p, x, l, mid, cur -> l);
        else update(p, x, mid + 1, r, cur -> r);

        cur -> sum = (cur -> l -> sum + cur -> r -> sum) % MOD;
    }

    void update(int L, int R, int x, int l, int r, Node *cur) {
        if (L <= l && r <= R) {
            cur -> lazy = (cur -> lazy + x) % MOD;
            cur -> sum = (cur -> sum + (LL)(x) * (r - l + 1)) % MOD;
            return;
        }

        pushdown(cur, r - l + 1);

        int mid = (l + r) >> 1;
        if (L <= mid) update(L, R, x, l, mid, cur -> l);
        if (mid < R) update(L, R, x, mid + 1, r, cur -> r);

        cur -> sum = (cur -> l -> sum + cur -> r -> sum) % MOD;
    }

    vector<int> bonus(int n, vector<vector<int>>& leadership,
                      vector<vector<int>>& operations) {
        vector<vector<int>> graph(n + 1);

        for (const auto &v : leadership)
            graph[v[0]].push_back(v[1]);

        vector<int> left(n + 1), right(n + 1);

        int ts = 0;
        dfs(1, ts, graph, left, right);


        Node *root = build(1, n);

        vector<int> ans;
        for (const auto &op : operations) {
            if (op[0] == 1) update(left[op[1]], op[2], 1, n, root);
            else if (op[0] == 2) update(left[op[1]], right[op[1]], op[2], 1, n, root);
            else ans.push_back(query(left[op[1]], right[op[1]], 1, n, root));
        }

        return ans;
    }
};