LCP 05 发 LeetCoin

Posted 2021-07-28 So istes immer

力扣决定给一个刷题团队发LeetCoin作为奖励。同时，为了监控给大家发了多少LeetCoin，力扣有时候也会进行查询。
该刷题团队的管理模式可以用一棵树表示：
团队只有一个负责人，编号为1。除了该负责人外，每个人有且仅有一个领导（负责人没有领导）；
不存在循环管理的情况，如A管理B，B管理C，C管理A。

力扣想进行的操作有以下三种：
给团队的一个成员（也可以是负责人）发一定数量的LeetCoin；
给团队的一个成员（也可以是负责人），以及他/她管理的所有人（即他/她的下属、他/她下属的下属，……），发一定数量的LeetCoin；
查询某一个成员（也可以是负责人），以及他/她管理的所有人被发到的LeetCoin之和。

输入：
N表示团队成员的个数（编号为1～N，负责人为1）；
leadership是大小为(N - 1) * 2的二维数组，其中每个元素[a, b]代表b是a的下属；
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)。

示例 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.

寻常方法：
最后一个测试样例过不了

class Node:
    def __init__(self):
        self.parent = None
        self.coin = 0
        self.child = []

class Solution(object):
    def __init__(self):
        self.g = collections.defaultdict(Node)

    def updateSingle(self, root, val):
        while root:
            root.coin += val   #找到指定结点，更新coin，并从下往上，更新在它之上的所有结点的coin
            root = root.parent

    def updateAll_2(self, root, val):
        res = val
        for i in root.child:            #用递归的方式，更新该结点和该结点下面所有结点的coin
            res += self.updateAll_2(i, val)
        root.coin += res
        return res

    def updateAll(self, root, val):
        res = self.updateAll_2(root,val)#得到指定结点更新后的coin
        root = root.parent              #从下往上，更新在它之上的所有结点的coin
        while root:
            root.coin += res
            root = root.parent

    def query(self, root):
        return root.coin % (10**9+7)

    def bonus(self, n, leadership, operations):
        for i,j in leadership: #遍历leadership数组，记录结点的父节点和子节点
            self.g[i].child.append(self.g[j])   
            self.g[j].parent = self.g[i] 
        res = []
        for op in operations:
            if op[0] == 1:     #第一种操作，更新单个结点
                self.updateSingle(self.g[op[1]],op[2]) 
            elif op[0] == 2:   #第二种操作，更新某个结点和它所有的孩子结点
                self.updateAll(self.g[op[1]],op[2])
            else:              #第三种操作，记录每次查询的结果
                res.append(self.query(self.g[op[1]]))
        return res