const MOD = 1000000007
type BitTree struct {
ar []int
n int
}
func NewBitTree(n int) *BitTree {
return &BitTree{make([]int, n+1), n}
}
func (bt *BitTree) Add(p int, v int) {
p++
for p <= bt.n {
bt.ar[p] += v
p += p & -p
}
}
func (bt *BitTree) Sum(p int) int {
o := 0
p++
for p > 0 {
o += bt.ar[p]
p -= p & -p
}
return o
}
func bonus(n int, leadership [][]int, operations [][]int) []int {
subs := make([][]int, n)
for _, l := range leadership {
a := l[0] - 1
subs[a] = append(subs[a], l[1]-1)
}
arIn := make([]int, n)
arOut := make([]int, n)
i := 0
var dfs func(int)
dfs = func(a int) {
arIn[a], i = i, i+1
for _, b := range subs[a] {
dfs(b)
}
arOut[a] = i - 1
}
dfs(0)
bt1, bt2 := NewBitTree(n), NewBitTree(n)
add := func(l, r, v int) {
bt1.Add(l, v)
bt1.Add(r+1, -v)
bt2.Add(l, v*(l-1))
bt2.Add(r+1, -v*r)
}
sum := func(l, r int) int {
l--
o := ((bt1.Sum(r)*r - bt2.Sum(r)) - (bt1.Sum(l)*l - bt2.Sum(l))) % MOD
if o < 0 {
o += MOD
}
return o
}
o := make([]int, 0, 32)
for _, op := range operations {
a := op[1] - 1
switch op[0] {
case 1:
add(arIn[a], arIn[a], op[2])
case 2:
add(arIn[a], arOut[a], op[2])
case 3:
o = append(o, sum(arIn[a], arOut[a]))
}
}
return o
}
M = int(1e9 + 7)
class BIT:
def __init__(self, n):
self.n = n + 5
self.sum = [0 for _ in range(n + 10)]
self.ntimessum = [0 for _ in range(n + 10)]
def lowbit(self, x):
return x & (-x)
# 在 pos 位置加上 k
def update(self, pos, k):
x = pos
while pos <= self.n:
self.sum[pos] += k
self.sum[pos] %= M
self.ntimessum[pos] += k * (x - 1)
self.ntimessum[pos] %= M
pos += self.lowbit(pos)
# 区间更新 + 单点查询
def askis(self, pos):
if not pos:
return 0
ret = 0
while pos:
ret += self.sum[pos]
ret %= M
pos -= lowbit(pos)
return ret
# 单点更新 + 区间查询
def asksi(self, l, r):
if l > r:
return 0
return askis(r) - askis(l - 1)
# 单点更新 + 单点查询
def askss(self, pos):
return askis(pos) - askis(pos - 1)
# 区间更新 + 区间查询
def askii(self, pos):
if not pos:
return 0
ret = 0
x = pos
while pos:
ret += x * self.sum[pos] - self.ntimessum[pos]
ret %= M
pos -= self.lowbit(pos)
return ret
class Solution:
def bonus(self, n: int, leadership: List[List[int]], operations: List[List[int]]) -> List[int]:
# 邻接表
g = [[] for _ in range(n + 1)]
begin = [0 for _ in range(n + 1)]
end = [0 for _ in range(n + 1)]
id = 1
for l in leadership:
g[l[0]].append(l[1])
# 深搜
def dfs(cur):
nonlocal id
begin[cur] = id
for child in g[cur]:
dfs(child)
end[cur] = id
id += 1
dfs(1)
# 树状数组
b = BIT(n)
ret = []
for q in operations:
if q[0] == 1:
b.update(end[q[1]], q[2])
b.update(end[q[1]] + 1, -q[2])
elif q[0] == 2:
b.update(begin[q[1]], q[2])
b.update(end[q[1]] + 1, -q[2])
else:
ans = b.askii(end[q[1]]) - b.askii(begin[q[1]] - 1)
ret.append((ans % M + M) % M)
return ret
'''
使用DFS序遍历子树，并保存相应的DFS序，DFS序可以保证任意子树的所有节点在一段连续空间上，
并保存节点编号到数组索引的映射
计算每一棵子树的节点数（等价于前一步骤中所有子树在数组上的连续空间长度）
后续就是标准的树状数组区间更新及区间查询写法（也可以使用线段树，但树状数组的常数更小），
单点更新可以替换为长度为1个区间更新
'''
mod = 10 ** 9 + 7
def lb(x):
return x & (-x)
class Solution:
def bonus(self, n: int, leadership: List[List[int]], operations: List[List[int]]) -> List[int]:
g1 = defaultdict(list)
g2 = [0] * (n + 1)
cnt = [0] + [1] * n
for a, b in leadership:
g1[a].append(b)
g2[b] = a
ts, stk = [0], [1]
while stk:
a = stk.pop()
ts.append(a)
stk.extend(g1[a][::-1])
d = {j: i for i, j in enumerate(ts)}
for b in ts[1:][::-1]:
cnt[g2[b]] += cnt[b]
maxn = n
while maxn & (maxn - 1):
maxn += lb(maxn)
f1 = [0] * (maxn + 1)
f2 = [0] * (maxn + 1)
def add(i, x):
xx = i * x
while i <= maxn:
f1[i] += x
f2[i] += xx
i += lb(i)
def query(i):
t1, t2, p = 0, 0, i
while i:
t1 += f1[i]
t2 += f2[i]
i -= lb(i)
return ((p + 1) * t1 - t2) % mod
ans = []
for op in operations:
if op[0] == 3:
p = op[1]
ans.append((query(d[p] + cnt[p] - 1) - query(d[p] - 1)) % mod)
else:
p, c, l = op[1], op[2], 1 if op[0] == 1 else cnt[op[1]]
add(d[p], c)
add(d[p] + l, -c)
return ans
