1534. 统计好三元组（中等）- LeetCode

题目描述

解法

题目的要求是求三元组个数，按照暴力解法的话算法复杂度会达到$O(n^3)$，一个比较简便的做法是先求二元组，比如：
$$|arr[j]-arr[k]|\le b$$
对于每一个二元组，满足三元组要求的第三个数满足:
$$|arr[i]-arr[j]|\le a$$
$$|arr[i]-arr[k]|\le c$$
拆开绝对值号之后，可以得到符合条件的值，一定落在$[arr[j]-a,arr[j]+a]$和$[arr[k]-c,arr[k]+c]$的交集，记为$[l,r]$。为了方便求取交集中数字的个数。我们新开一个记录数字出现频次的数组nums,对于落在区间$[l,r]$的数字的个数，可以得到答案$nums[r]-nums[l-1]$。最后算法的复杂度为$O(n^2)$。

class Solution:
    def countGoodTriplets(self, arr: List[int], a: int, b: int, c: int) -> int:
        nums = [0] * 1001
        ans = 0
        length = len(arr)
        for j in range(length):
            for k in range(j+1,length):
                if abs(arr[j]-arr[k]) <= b:
                    r = min(arr[j]+a,arr[k]+c,1000)
                    l = max(arr[j]-a,arr[k]-c,0)
                    if r >= l:
                        ans += nums[r] if l == 0 else nums[r] - nums[l-1]
            for m in range(arr[j],1001):
                nums[m] += 1
        return ans 

官方解答参考链接：官方解答。