LCP 05. 发 LeetCoin

题目描述

力扣决定给一个刷题团队发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的二维数组，代表以时间排序的操作，格式如下：
   1. operations[i][0] = 1: 代表第一种操作，operations[i][1]代表成员的编号，operations[i][2]代表LeetCoin的数量；
   2. operations[i][0] = 2: 代表第二种操作，operations[i][1]代表成员的编号，operations[i][2]代表LeetCoin的数量；
   3. operations[i][0] = 3: 代表第三种操作，operations[i][1]代表成员的编号；

输出：

返回一个数组，数组里是每次查询的返回值（发LeetCoin的操作不需要任何返回值）。由于发的LeetCoin很多，请把每次查询的结果模1e9+7 (1000000007)。

 

示例 1：

输入：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

解法

方法一：线段树

线段树将整个区间分割为多个不连续的子区间，子区间的数量不超过 log(width)。更新某个元素的值，只需要更新 log(width) 个区间，并且这些区间都包含在一个包含该元素的大区间内。区间修改时，需要使用懒标记保证效率。

• 线段树的每个节点代表一个区间；
• 线段树具有唯一的根节点，代表的区间是整个统计范围，如 [1, N]；
• 线段树的每个叶子节点代表一个长度为 1 的元区间 [x, x]；
• 对于每个内部节点 [l, r]，它的左儿子是 [l, mid]，右儿子是 [mid + 1, r], 其中 mid = ⌊(l + r) / 2⌋ (即向下取整)。

Python3

MOD = int(1e9 + 7)


class Node:
    def __init__(self, l, r):
        self.left = None
        self.right = None
        self.l = l
        self.r = r
        self.mid = (l + r) >> 1
        self.v = 0
        self.add = 0


class SegmentTree:
    def __init__(self, n):
        self.root = Node(1, n)

    def modify(self, l, r, v, node=None):
        if l > r:
            return
        if node is None:
            node = self.root
        if node.l >= l and node.r <= r:
            node.v = (node.v + (node.r - node.l + 1) * v) % MOD
            node.add += v
            return
        self.pushdown(node)
        if l <= node.mid:
            self.modify(l, r, v, node.left)
        if r > node.mid:
            self.modify(l, r, v, node.right)
        self.pushup(node)

    def query(self, l, r, node=None):
        if l > r:
            return 0
        if node is None:
            node = self.root
        if node.l >= l and node.r <= r:
            return node.v
        self.pushdown(node)
        v = 0
        if l <= node.mid:
            v += self.query(l, r, node.left)
        if r > node.mid:
            v += self.query(l, r, node.right)
        return v % MOD

    def pushup(self, node):
        node.v = (node.left.v + node.right.v) % MOD

    def pushdown(self, node):
        if node.left is None:
            node.left = Node(node.l, node.mid)
        if node.right is None:
            node.right = Node(node.mid + 1, node.r)
        if node.add:
            left, right = node.left, node.right
            left.v = (left.v + (left.r - left.l + 1) * node.add) % MOD
            right.v = (right.v + (right.r - right.l + 1) * node.add) % MOD
            left.add += node.add
            right.add += node.add
            node.add = 0


class Solution:
    def bonus(
        self, n: int, leadership: List[List[int]], operations: List[List[int]]
    ) -> List[int]:
        def dfs(u):
            nonlocal idx
            begin[u] = idx
            for v in g[u]:
                dfs(v)
            end[u] = idx
            idx += 1

        g = defaultdict(list)
        for a, b in leadership:
            g[a].append(b)
        begin = [0] * (n + 1)
        end = [0] * (n + 1)
        idx = 1
        dfs(1)
        ans = []
        tree = SegmentTree(n)
        for op in operations:
            p, v = op[:2]
            if p == 1:
                tree.modify(end[v], end[v], op[2])
            elif p == 2:
                tree.modify(begin[v], end[v], op[2])
            else:
                ans.append(tree.query(begin[v], end[v]))
        return ans

Java

class Node {
    Node left;
    Node right;
    int l;
    int r;
    int mid;
    int v;
    int add;
    public Node(int l, int r) {
        this.l = l;
        this.r = r;
        this.mid = (l + r) >> 1;
    }
}

class SegmentTree {
    private Node root;
    private static final int MOD = (int) 1e9 + 7;

    public SegmentTree(int n) {
        root = new Node(1, n);
    }

    public void modify(int l, int r, int v) {
        modify(l, r, v, root);
    }

    public void modify(int l, int r, int v, Node node) {
        if (l > r) {
            return;
        }
        if (node.l >= l && node.r <= r) {
            node.v = (node.v + (node.r - node.l + 1) * v) % MOD;
            node.add += v;
            return;
        }
        pushdown(node);
        if (l <= node.mid) {
            modify(l, r, v, node.left);
        }
        if (r > node.mid) {
            modify(l, r, v, node.right);
        }
        pushup(node);
    }

    public int query(int l, int r) {
        return query(l, r, root);
    }

    public int query(int l, int r, Node node) {
        if (l > r) {
            return 0;
        }
        if (node.l >= l && node.r <= r) {
            return node.v;
        }
        pushdown(node);
        int v = 0;
        if (l <= node.mid) {
            v = (v + query(l, r, node.left)) % MOD;
        }
        if (r > node.mid) {
            v = (v + query(l, r, node.right)) % MOD;
        }
        return v;
    }

    public void pushup(Node node) {
        node.v = (node.left.v + node.right.v) % MOD;
    }

    public void pushdown(Node node) {
        if (node.left == null) {
            node.left = new Node(node.l, node.mid);
        }
        if (node.right == null) {
            node.right = new Node(node.mid + 1, node.r);
        }
        if (node.add != 0) {
            Node left = node.left, right = node.right;
            left.v = (left.v + (left.r - left.l + 1) * node.add) % MOD;
            right.v = (right.v + (right.r - right.l + 1) * node.add) % MOD;
            left.add += node.add;
            right.add += node.add;
            node.add = 0;
        }
    }
}

class Solution {
    private List<Integer>[] g;
    private int[] begin;
    private int[] end;
    private int idx;

    public int[] bonus(int n, int[][] leadership, int[][] operations) {
        g = new List[n + 1];
        Arrays.setAll(g, k -> new ArrayList<>());
        for (int[] l : leadership) {
            int a = l[0], b = l[1];
            g[a].add(b);
        }
        begin = new int[n + 1];
        end = new int[n + 1];
        idx = 1;
        dfs(1);
        List<Integer> ans = new ArrayList<>();
        SegmentTree tree = new SegmentTree(n);
        for (int[] op : operations) {
            int p = op[0], v = op[1];
            if (p == 1) {
                tree.modify(end[v], end[v], op[2]);
            } else if (p == 2) {
                tree.modify(begin[v], end[v], op[2]);
            } else {
                ans.add(tree.query(begin[v], end[v]));
            }
        }
        return ans.stream().mapToInt(Integer::intValue).toArray();
    }

    private void dfs(int u) {
        begin[u] = idx;
        for (int v : g[u]) {
            dfs(v);
        }
        end[u] = idx;
        ++idx;
    }
}

C++

const int MOD = 1e9 + 7;

class Node {
public:
    Node* left;
    Node* right;
    int l;
    int r;
    int mid;
    int v;
    int add;

    Node(int l, int r) {
        this->l = l;
        this->r = r;
        this->mid = (l + r) >> 1;
        this->left = this->right = nullptr;
        v = add = 0;
    }
};

class SegmentTree {
private:
    Node* root;

public:
    SegmentTree(int n) {
        root = new Node(1, n);
    }

    void modify(int l, int r, int v) {
        modify(l, r, v, root);
    }

    void modify(int l, int r, int v, Node* node) {
        if (l > r) return;
        if (node->l >= l && node->r <= r) {
            node->v = (node->v + (node->r - node->l + 1) * v) % MOD;
            node->add += v;
            return;
        }
        pushdown(node);
        if (l <= node->mid) modify(l, r, v, node->left);
        if (r > node->mid) modify(l, r, v, node->right);
        pushup(node);
    }

    int query(int l, int r) {
        return query(l, r, root);
    }

    int query(int l, int r, Node* node) {
        if (l > r) return 0;
        if (node->l >= l && node->r <= r) return node->v;
        pushdown(node);
        int v = 0;
        if (l <= node->mid) v += query(l, r, node->left);
        if (r > node->mid) v += query(l, r, node->right);
        return v % MOD;
    }

    void pushup(Node* node) {
        node->v = (node->left->v + node->right->v) % MOD;
    }

    void pushdown(Node* node) {
        if (!node->left) node->left = new Node(node->l, node->mid);
        if (!node->right) node->right = new Node(node->mid + 1, node->r);
        if (node->add) {
            Node* left = node->left;
            Node* right = node->right;
            left->v = (left->v + (left->r - left->l + 1) * node->add) % MOD;
            right->v = (right->v + (right->r - right->l + 1) * node->add) % MOD;
            left->add += node->add;
            right->add += node->add;
            node->add = 0;
        }
    }
};

class Solution {
public:
    int idx;

    vector<int> bonus(int n, vector<vector<int>>& leadership, vector<vector<int>>& operations) {
        vector<vector<int>> g(n + 1);
        for (auto& l : leadership) {
            int a = l[0], b = l[1];
            g[a].push_back(b);
        }
        vector<int> begin(n + 1);
        vector<int> end(n + 1);
        idx = 1;
        dfs(1, begin, end, g);
        vector<int> ans;
        SegmentTree* tree = new SegmentTree(n);
        for (auto& op : operations) {
            int p = op[0], v = op[1];
            if (p == 1)
                tree->modify(end[v], end[v], op[2]);
            else if (p == 2)
                tree->modify(begin[v], end[v], op[2]);
            else
                ans.push_back(tree->query(begin[v], end[v]));
        }
        return ans;
    }

    void dfs(int u, vector<int>& begin, vector<int>& end, vector<vector<int>>& g) {
        begin[u] = idx;
        for (int v : g[u]) dfs(v, begin, end, g);
        end[u] = idx;
        ++idx;
    }
};

...