[Leetcode] LCP 05. 发LeetCoin

题目链接:https://leetcode-cn.com/problems/coin-bonus/
分析:
直接dfs去求和会超时。
利用线段树,进行区间求和和更新。
比较难理解的点是将团队成员id变成连续的。将求解一个子集合的和问题转为求解一个连续区间的和问题。对这个连续区间就可以二分去求了。
所以要先使用dfs对原来的id进行映射。之后使用的坐标就是区间在不断二分的二叉树中的位置。

主要操作是update和query。更新父结点时,可以先不继续更新子结点。而是保存到lazy数组中,直到访问子结点时再往下传递,即需要一个pushdown的操作。
Python

class Solution:
    cnt = 0
    def bonus(self, n: int, leadership: List[List[int]], operations: List[List[int]]) -> List[int]:
        mod = 10**9+7
        rangeSum = [0]*(4*n)
        lazy = [0]*(4*n)
        left = [0]*(n+1)
        right = [0]*(n+1)
        dic = {}
        self.cnt = 0 
        for i in range(1, n+1):
            dic[i] = []
        def update(pos, left, right, opLeft, opRight, k):
            if (opLeft > right or opRight < left):
                return
            if (opLeft <= left and opRight >= right):
                rangeSum[pos] = (rangeSum[pos]+(right-left+1)*k) % mod
                lazy[pos] = (lazy[pos] + k) % mod
                return 
            pushDown(pos, left, right)
            mid = (left+right)>>1
            update(pos*2, left, mid, opLeft, opRight, k)
            update(pos*2+1, mid+1, right, opLeft, opRight, k)
            rangeSum[pos] = (rangeSum[pos*2] + rangeSum[pos*2+1])%mod

        def pushDown(pos, left, right):
            if lazy[pos] == 0:
                return 
            mid = (left+right)>>1
            rangeSum[pos*2] = (rangeSum[pos*2]+lazy[pos]*(mid-left+1))%mod
            rangeSum[pos*2+1] = (rangeSum[pos*2+1]+lazy[pos]*(right-mid))%mod
            lazy[pos*2] = (lazy[pos*2]+lazy[pos])%mod
            lazy[pos*2+1] = (lazy[pos*2+1]+lazy[pos])%mod
            lazy[pos] = 0

        def query(pos, left, right, qLeft, qRight):
            if (qLeft > right or qRight < left):
                return 0
            if (qLeft <= left and qRight >= right):
                return rangeSum[pos]
            pushDown(pos, left, right)
            mid = (left + right)>>1
            return (query(pos*2, left, mid, qLeft, qRight) 
            + query(pos*2+1, mid+1, right, qLeft, qRight))%mod

        
        def dfs(u):
            self.cnt += 1
            left[u] = self.cnt
            
            for i in dic[u]:
                dfs(i)
            right[u] = self.cnt

        res = []
        for i in leadership:
            dic[i[0]].append(i[1])
        dfs(1)        

        print(left)
        for operation in operations:
            if operation[0] == 1:
                update(1, 1, n, left[operation[1]], left[operation[1]], operation[2])
            elif operation[0] == 2:
                update(1, 1, n, left[operation[1]], right[operation[1]], operation[2])
            elif operation[0] == 3:
                res.append(query(1, 1, n, left[operation[1]], right[operation[1]]))
        return res
posted @ 2020-10-16 15:11  我的小叮当  阅读(134)  评论(0编辑  收藏  举报