1534. 统计好三元组

1. 三重循环O(N^3)，O(1)
2. 首先用一个双重循环筛选出满足|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