Question

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⌋ (即向下取整)。

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

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;
  }
}

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;
  }
};