day20 统计好三元组(leetcode)

1534. 统计好三元组

题目描述

给你一个整数数组 arr ，以及 a、b 、c 三个整数。请你统计其中好三元组的数量。

如果三元组 (arr[i], arr[j], arr[k]) 满足下列全部条件，则认为它是一个 好三元组 。

0 <= i < j < k < arr.length
|arr[i] - arr[j]| <= a
|arr[j] - arr[k]| <= b
|arr[i] - arr[k]| <= c

其中 |x| 表示 x 的绝对值。

返回 好三元组的数量 。

解题：

法一：暴力破解

写三个循环，用$O(n^3)$的时间复杂度解决问题

class Solution:
    def countGoodTriplets(self, arr: List[int], a: int, b: int, c: int) -> int:
        n = len(arr)
        cnt = 0
        for i in range(n):
            for j in range(i + 1, n):
                for k in range(j + 1, n):
                    if abs(arr[i] - arr[j]) <= a and abs(arr[j] - arr[k]) <= b and abs(arr[i] - arr[k]) <= c:
                        cnt += 1
        return cnt

法二：前缀和优化

1. 首先，我们仍然要用$O(n^2)$ 的时间复杂度来找到符合条件的每一个j,k
2. 其次，根据选定的j,k，我们可以计算出符合条件的arr[i]的值域
3. 然后，我们考虑到用符合条件的i的前缀和来优化（我们需要得到满足条件的i的个数）
4. 前缀和数组每次更新（这样确保i永远是小于等于）

class Solution:
    def countGoodTriplets(self, arr: List[int], a: int, b: int, c: int) -> int:
        ans = 0
        n = len(arr)
        total = [0] * 1001
        for j in range(n):
            for k in range(j + 1, n):
                if abs(arr[j] - arr[k]) <= b:
                	# 计算值域
                    lj, rj = arr[j] - a, arr[j] + a
                    lk, rk = arr[k] - c, arr[k] + c
                    l = max(0, lj, lk)
                    r = min(1000, rj, rk)
                    # 计算词频
                    if l <= r:
                        ans += total[r] if l == 0 else total[r] - total[l - 1]
            # 更新前缀和数组
            for k in range(arr[j], 1001):
                total[k] += 1
        
        return ans