跳转至

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⌋ (即向下取整)。
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
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

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
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;
    }
}

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
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;
    }
};

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
class Node {
    var left: Node?
    var right: Node?
    let l: Int
    let r: Int
    let mid: Int
    var v = 0
    var add = 0

    init(_ l: Int, _ r: Int) {
        self.l = l
        self.r = r
        self.mid = (l + r) >> 1
    }
}

class SegmentTree {
    private var root: Node
    private let MOD = 1_000_000_007

    init(_ n: Int) {
        root = Node(1, n)
    }

    func modify(_ l: Int, _ r: Int, _ v: Int) {
        modify(l, r, v, root)
    }

    private func modify(_ l: Int, _ r: Int, _ v: Int, _ 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 = (node.add + v) % MOD
            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)
    }

    func query(_ l: Int, _ r: Int) -> Int {
        return query(l, r, root)
    }

    private func query(_ l: Int, _ r: Int, _ node: Node) -> Int {
        if l > r {
            return 0
        }
        if node.l >= l && node.r <= r {
            return node.v
        }
        pushdown(node)
        var 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
    }

    private func pushup(_ node: Node) {
        node.v = (node.left!.v + node.right!.v) % MOD
    }

    private func pushdown(_ node: Node) {
        if node.left == nil {
            node.left = Node(node.l, node.mid)
        }
        if node.right == nil {
            node.right = Node(node.mid + 1, node.r)
        }
        if node.add != 0 {
            let 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 = (left.add + node.add) % MOD
            right.add = (right.add + node.add) % MOD
            node.add = 0
        }
    }
}

class Solution {
    private var g: [[Int]] = []
    private var begin: [Int] = []
    private var end: [Int] = []
    private var idx = 1

    func bonus(_ n: Int, _ leadership: [[Int]], _ operations: [[Int]]) -> [Int] {
        g = Array(repeating: [], count: n + 1)
        for l in leadership {
            let (a, b) = (l[0], l[1])
            g[a].append(b)
        }

        begin = Array(repeating: 0, count: n + 1)
        end = Array(repeating: 0, count: n + 1)
        idx = 1
        dfs(1)

        var ans: [Int] = []
        let tree = SegmentTree(n)
        for op in operations {
            let (p, v) = (op[0], op[1])
            if p == 1 {
                tree.modify(begin[v], begin[v], op[2])
            } else if p == 2 {
                tree.modify(begin[v], end[v], op[2])
            } else if p == 3 {
                ans.append(tree.query(begin[v], end[v]))
            }
        }
        return ans
    }

    private func dfs(_ u: Int) {
        begin[u] = idx
        for v in g[u] {
            dfs(v)
        }
        end[u] = idx
        idx += 1
    }
}

评论