1534. 统计好三元组

三重循环O(N^3)，O(1)首先用一个双重循环筛选出满足|arr[j] - arr[k]| <= b的 j 和 k，再从中筛选出满足如下两个条件的 i：i<j ，|arr[i] - arr[j]| <= a，|arr[i] - arr[k]| <= c。而此时不必再遍历数组完成查找。我们可以维护一个 arr[i] 频次数组的前缀和 pre，即pre[i]=k表示 arr 数组中小于等于 i 的数有 k 个，满足第二个条件的数的个数即为pre[rb]-pre[lb-1]。又有，我们在每个 j 循环的末尾才更新 pre 数组，这样就保证我们每次是在小于 j 的数中查找的。O(N^2+N*S)，O(S) class Solution: def countGoodTriplets(self, arr: List[int], a: int, b: int, c: int) -> int: ans = 0 l = len(arr) pre = [0]*1001 #pre[i]=k 表示小于等于i的元素有k个 for j in range(l): for k in range(j+1,l): if abs(arr[j]-arr[k])<=b: lb = max(0,arr[j]-a,arr[k]-c) rb = min(1000,arr[j]+a,arr[k]+c) if rb>=lb: ans+=pre[rb]-pre[lb-1] if lb>0 else pre[rb] for i in range(arr[j],1001): pre[i]+=1 return ans

当S比较大的时候，可以采用树状数组。

class BIT: def __init__(self, n): self.n = n # n为原数组的长度 self.tree = [0] * (n + 1) @staticmethod def lowbit(x): return x & (-x) def query(self, idx): ret = 0 while idx > 0: ret += self.tree[idx] idx -= BIT.lowbit(idx) return ret def update(self, idx, v): # v是更新前后的插值 while idx <= self.n: self.tree[idx] += v idx += BIT.lowbit(idx) class Solution: def countGoodTriplets(self, arr: List[int], a: int, b: int, c: int) -> int: ans = 0 l = len(arr) bit = BIT(1002) for j in range(l): for k in range(j+1,l): if abs(arr[j]-arr[k])<=b: lb = max(0,arr[j]-a,arr[k]-c) rb = min(1001,arr[j]+a,arr[k]+c) if rb>=lb: ans+=bit.query(rb+1)-bit.query(lb) if lb>0 else bit.query(rb+1) bit.update(arr[j]+1,1) return ans