leetcode 2916. 子数组不同元素数目的平方和 II（区间更新 + 区间查询 线段树第二个板子 双闭区间 避开0）

区间更新 + 区间查询 线段树第二个板子 双闭区间 避开0）"/>

leetcode 2916. 子数组不同元素数目的平方和 II（区间更新 + 区间查询 线段树第二个板子 双闭区间 避开0）

描述

偷了一个线段树板子
不知道为啥要避开0
然后这里的更新和查找都是用双闭区间的

ac code

class SegmentTree:def __init__(self, n):self.n = n self.B1 = [0]*n self.B2 = [0]*n  def add(self, b, idx, x):N = self.n while idx < N:b[idx] += xidx += idx & -idxdef range_add(self, l,r,x):B1 = self.B1B2 = self.B2self.add(B1, l, x)self.add(B1, r+1, -x)self.add(B2, l, x*(l-1))self.add(B2, r+1, -x*r)def sum(self, b, idx):total = 0while idx > 0:total += b[idx]idx -= idx & -idxreturn totaldef prefix_sum(self, idx):B1 = self.B1B2 = self.B2return self.sum(B1, idx)*idx -  self.sum(B2, idx)def range_sum(self, l, r):return self.prefix_sum(r) - self.prefix_sum(l-1)class Solution:def sumCounts(self, nums: List[int]) -> int:n = len(nums)seg = SegmentTree(n+5)MOD = 10**9+7seen = dict()ans, cur = 0, 0 for i, a in enumerate(nums):j = -1 if a not in seen else seen[a]s = seg.range_sum(j+2,i+1) # 双闭区间cur += s*2+i-jcur %= MOD ans += curans %= MOD# print(i,s,cur)seg.range_add(j+2,i+1,1)seen[a] = ireturn ans

题目链接

lc2916