[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