[前缀和, 枚举中间] 统计好三元组

题目

题目链接: https://leetcode.cn/problems/count-good-triplets/description/

给你一个整数数组 \(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\)

方法一: 暴力枚举

复杂度分析

  • 时间复杂度 \(O(n^3)\), n 是 \(arr\) 数组的长度;
  • 空间复杂度 \(O(1)\).

方法二: 动态维护、前缀和

考虑枚举 \(j, k\), 定义 \(A_i = arr[i], A_j = arr[j], A_k = arr[k]\).

那么对于 \(i < j < k\) 满足的 \(i\) 如下条件:

  • \(A_j - a ≤ A_i ≤ A_j + a\)
  • \(A_k - c ≤ A_i ≤ A_k + c\)

合并两个区间我们能够得到 \(A_i\) 的范围如下:

\[[max(0, A_j - a, A_k - c), min(A_j + a, A_k + c, mx)] \]

本题也就转换成一个单点修改、区间查询的问题。对于单点修改和区间查询，我们可以考虑前缀和、树状数组和线段树。

  • 单点修改: 随着 \(j\) 的枚举, 我们可以先查询再修改, 实现一个动态维护, 可以避免 \(i\) 维度的枚举。
  • 区间查询: 计算 \(A_i\) 范围, 边界条件即为查询的区间。

代码

前缀和 
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;

#define len(f) (int)f.size()
#define mst(f, x) memset(f, x, sizeof(f));
const int inf = 0x3f3f3f3f, MOD = 1e9 + 7;

int countGoodTriplets(vector<int> &arr, int a, int b, int c) {
    int n = len(arr), mx = ranges::max(arr);
    int ans = 0;

    vector<int> pre(mx + 2);
    for (int j = 0; j < n; j++) {
      int y = arr[j];
      // 先更新
      for (int k = j + 1; k < n; ++k) {
        int z = arr[k];
        if (abs(y - z) > b) {
          continue;
        }
        int l = max({y - a, z - c, 0});
        int r = min({y + a, z + c, mx});
        ans += max(pre[r + 1] - pre[l], 0);
      }
      // 后动态维护
      for (int v = y + 1; v < mx + 2; v++) {
        pre[v]++; // 将 y 加入 cnt 数组, 维护所有受影响的前缀和
      }
    }
    return ans;
}

复杂度分析

  • 时间复杂度 \(O(n(n + U))\), U 是 arr 的最大值;
  • 空间复杂度 \(O(U)\).

相关题型

方法三: 排序、枚举中间、三指针/二分

如果 \(A_i ≤ 1e9\)，那么方法二的 \(O(n(n+U))\) 的时间复杂度会超时，下面介绍一种时间复杂度与值域 \(U\) 无关的方法。

❌直接排序：题目要求 \(i < j < k\)，不能打乱顺序。
✅下标排序：对下标数组按 \(A_i\) 的值从小到大排序。

vector<int> id(n);
iota(id.begin(), id.end(), 0);
sort(id.begin(), id.end(), {}, [&](int i) { return arr[i]; });

遍历 id 的元素作为下标 \(j\)。（枚举中间）

  • 把满足 \(i < j\) 且 \(|A_i - A_j| ≤ a\) 的 \(A_i\) 保存到数组 \(left\) 中。
  • 把满足 \(j < k\) 且 \(|A_k - A_j| ≤ b\) 的 \(A_k\) 保存到数组 \(right\) 中。

我们就能够得到两个有序数组，接下来我们就把问题变成：

  • 从两个有序数组中，分别选一个数，计算满足绝对差 \(≤ c\) 的数对个数。

我们遍历 \(left\) 中的元素 \(x\)，计算 \(right\) 中有多少个元素 \(z\) 满足 \(|x - z| ≤ c\) 。即 \([x - c, x + c]\) 中 \(z\) 的个数。

  • \(right\) 中的 \(≤x+c\) 的元素个数，减去 \(right\) 中的 \(<x−c\) 的元素个数。

两种解法：

  • 三指针
int k1 = 0, k2 = 0;
for (int x : left) {
	// - 三指针
	while (k2 < len(right) && right[k2] <= x + c) k2++;
	while (k1 < len(right) && right[k1] < x - c) k1++;
	ans += k2 - k1;
}
  • 二分
for (int x : left) {
    int cnt = ranges::upper_bound(right, x+c) - ranges::lower_bound(right, x-c);
    ans += cnt;
}

代码

int countGoodTriplets(vector<int> &arr, int a, int b, int c) {
    int n = len(arr);
    // 创建下标数组 id, 按 arr[i] 的值进行非递减排序
    vector<int> id(n);
    ranges::iota(id, 0);
    ranges::sort(id, {}, [&](int i) { return arr[i]; });

    // 创建两个有序数组 left 和 right, 枚举 id 作为 j
    // - (1) left 中的元素满足 i < j 且 |arr[i] - arr[j| ≤ a
    // - (2) right 中的元素满足 j < k 且 |arr[k] - arr[j| ≤ b
    int ans = 0;
    for (int j : id) {
      int y = arr[j];
      vector<int> left, right;
      for (int i : id) {
        if (i < j && abs(arr[i] - y) <= a) {
          left.push_back(arr[i]);
        }
      }
      for (int k : id) {
        if (j < k && abs(arr[k] - y) <= b) {
          right.push_back(arr[k]);
        }
      }

      // 给定两个有序数组，从两个数组中各选一个数，计算绝对差 ≤c 的数对个数。
      for (int x : left) {
        int cnt = ranges::upper_bound(right, x+c) - ranges::lower_bound(right, x-c);
        ans += cnt;
      }
    }
    return ans;
  }

复杂度分析

  • 时间复杂度 \(O(n^2)\);
  • 空间复杂度 \(O(n)\).
posted @ 2025-04-14 11:30  kris4js  阅读(37)  评论(0)    收藏  举报