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 的绝对值。

返回 好三元组的数量 。

其中：
3 <= arr.length <= 100
0 <= arr[i] <= 1000
0 <= a, b, c <= 1000

法一：暴力法

class Solution {
public:
    int countGoodTriplets(vector<int>& arr, int a, int b, int c) {
        vector<int>::size_type sz = arr.size();
        int cnt = 0;
        for (int i = 0; i < sz - 2; ++i) {
            for (int j = i + 1; j < sz - 1; ++j) {
                if (abs(arr[i] - arr[j]) <= a) {
                    for (int k = j + 1; k < sz; ++k) {
                        if (abs(arr[i] - arr[k]) <= c && abs(arr[j] - arr[k]) <= b) {
                            ++cnt;
                        }
                    }
                }
            }
        }
        return cnt;
    }
};

该方法的时间复杂度为O（n³），空间复杂度O（1）。

法二：枚举优化。我们先满足|arr[j] - arr[k]| <= b这一项，之后还要满足的情况为|arr[i] - arr[j]| <= a和|arr[i] - arr[k]| <= c，要满足的两个条件可去掉绝对值即arr[j] - a <= arr[i] <= arr[j] + a和arr[k] - c <= arr[i] <= arr[k] + c，将这两个条件取交集[l, r]，即要满足的两个条件可转化为此条件：arr[i]要在集合[l, r]中。我们可以创建arr[i]出现的频次前缀和数组，即下标为arr[i]可能取得的值、值为小于等于下标值的arr[i]的出现次数的数组。还有就是要满足i < j，这可以通过在频次前缀和数组中对arr[j]进行计数，arr[i]和arr[j]的出现频次前缀和数组都是相同的，以j为下标更新前缀和数组相当于实现了i < j：

class Solution {
public:
    int countGoodTriplets(vector<int>& arr, int a, int b, int c) {
        int cnt = 0;
        vector<int>::size_type sz = arr.size(); 
        vector<int> prefixSum(1001);

        for (int j = 0; j < sz; ++j) {
            for (int k = j + 1; k < sz; ++k) {    // 此处满足j < k
                if (abs(arr[j] - arr[k]) <= b) {
                    int lj = arr[j] - a, rj = arr[j] + a;
                    int lk = arr[k] - c, rk = arr[k] + c;
                    int l = max(0, max(lj, lk)), r = min(min(rj, rk), 1000);
                    
                    if (l <= r) {
                        if (l == 0) {
                            cnt += prefixSum[r];
                        } else {
                            cnt += prefixSum[r] - prefixSum[l - 1];
                        }
                    }
                }
            }
            for (int k = arr[j]; k <= 1000; ++k) {
                ++prefixSum[k];
            }
        }
        return cnt;
    }
};

该方法时间复杂度为O（n² + nS)，n为数组长度，S为输入数组值的值域，本例中为1001，而空间复杂度为O（S）。