[排列置换, 枚举中间, 树状数组, 前后缀分解] 统计数组中好三元组数目

题目

https://leetcode.cn/problems/count-good-triplets-in-an-array/

给你两个 \([0, n - 1]\) 的 排列 \(a, b\) 。\((∀i,j∈[0, n - 1],\ i ≠ j), a_i ≠ a_j\)。

定义「好三元组」：

• \(∀x, y, z ∈ [0, n - 1]\)
• 满足 \(a_x < a_y < a_z\) 且 \(b_x < b_y < b_z\) 。

返回排列中「好三元组」的数目。

• \(n = len(a) = len(b)\)
• \(3 ≤ n ≤ 1e5\)
• \(0 ≤ a_i, b_i ≤ n - 1\)

---

方法一：置换 + 枚举中间 + 树状数组 + 前后缀分解

题目的本质是求 \(a\) 和 \(b\) 的恰好长度为 3 的公共子序列的个数。

• 1143. 最长公共子序列。

但本题的 \(n\) 太大，对于 \(O(n^2)\) 的 DP 会爆时间复杂度。

思路

排列数组中的各个元素都不相同，如果我们能够把 \(a\) 变成 \([0, 1, 2, ...,n - 1]\)，那么我们就能够把「公共子序列问题」转换成拥有更好性质的「严格递增子序列问题」。

对于 3 个数的问题，通常可以枚举中间的元素。

1. 置换

置换是一个排列到另一个排列的双射。[排列，双射]

为了得到单调递增的排列 \(A = [0, 1, 2, ..., n - 1]\)，我们可以使用以下的方法：

vector<int> f(n);
for (int i = 0; i < n; i++) {
    f[a[i]] = i;
}

对于 \(b\) 的所有元素，如果采用了双射变换，得到的新排列 \(A, B\) 仍然满足 求长度为 3 的公共子序列的个数 问题。同时由于 \(A\) 的递增性质，我们只需要考虑 \(B\)。

• 新问题：\(B\) 有多少个长为 \(3\) 的递增子序列？

2. 枚举中间

对于长度为 \(3\) 的严格递增子序列 \((x, y, z)\)，我们考虑枚举中间元素 \(y\)。

• 新问题：在 \(B\) 中，元素 \(y\) 的左侧有多少个比 \(y\) 小的元素 \(x\)？右侧有多少个比 \(y\) 大的数 \(z\)？

3. 树状数组

类似于「最长递增子序列问题」。我们用值域树状数组（或有序集合 \(set\)）实现对左/右侧元素的维护。这里采用值域树状数组：

• 把元素视为下标 \(i\)，添加一个值为 \(3\) 的数，就是调用 \(add(3, 1)\)。
• 查询小于 \(3\) 的元素个数，即小于等于 \(2\) 的元素个数，就是调用 \(query(2)\)。

4.1 前后缀分解

在 枚举中间 步骤，我们将问题拆分为「在 \(B\) 中，元素 \(y\) 的左侧有多少个比 \(y\) 小的元素 \(x\)？右侧有多少个比 \(y\) 大的数 \(z\)？」。为了解决问题，我们可以考虑前后缀分解的方法，即先计算后缀 \(suf[i]\)，这样我们在就能够在枚举 \(y\) 的过程中获取 \(y\) 后面元素的信息，将 \(O(n^2)\) 优化为 \(O(2n)\)。

代码

template <typename T>
class FenWickTree {
public:
  vector<T> tree;
  int n;
  FenWickTree(int n): n(n), tree(n) {}

  int lowbit(int x) {
    return x & -x;
  }

  void add(int i, T x) {
    for (; i < n; i += lowbit(i)) {
      tree[i] += x;
    }
  }

  T query(int i) const {
    T res = 0;
    for (; i > 0; i &= i - 1) {
      res += tree[i];
    }
    return res;
  }
};

long long goodTriplets(vector<int> &a, vector<int> &b) {
  int n = len(a);
  vector<int> pos(n);
  for (int i = 0; i < n; i++) {
	pos[a[i]] = i; // a[i] -> pos[a[i]]
  }
	
  for (int i = 0; i < n; i++) {
	b[i] = pos[b[i]];
  }
	
  // [0, n-1] 的排列在树状数组中映射为 [1, n]
  FenWickTree<int> bit1(n + 1);
	
  vector<int> suf(n + 1);
  for (int i = n - 1; i >= 0; i--) {
	int x = b[i];
	suf[i] = (n - 1 - i) - bit1.query(x + 1);
	bit1.add(x + 1, 1);
  }
	
  ll ans = 0; // 记录答案
	
  FenWickTree<int> bit2(n + 1);
  int pre = 0;
  for (int i = 0; i < n; i++) {
	int x = b[i];
	pre = bit2.query(x);
	ans += (ll)pre * suf[i];
	bit2.add(x + 1, 1);
  }
  return ans;
}

复杂度

• 时间复杂度：\(O(nlogn)\)
• 空间复杂度：\(O(n)\)

4.2 排列性质优化

在前后缀分解中，我们的前缀和后缀分别记录了 前面有多少个小于 y 的元素 和 后面有多少个大于 y 的元素。

但实际上，由于排列的性质（所有元素都不相同），我们只需要知道前面有多少个小于 y 的元素，就能够计算出后面有多少个大于 y 的元素。

不妨假设当前枚举 \(i, B[i] = y\)，前面有 \(lo\) 个小于 \(y\) 的元素，那么有如下计算：

• 前面总共 \(i\) 个元素，那么大于 \(y\) 的有 \(i - lo\) 个元素；
• 整个排列大于 \(y\) 的一共有 \(n - 1 - y\)；
• 那么后面大于 \(y\) 的一共有 \(n - 1 - y - (i - lo)\) 个元素。

优化如下：

long long goodTriplets(vector<int> &a, vector<int> &b) {
  int n = len(a);
  vector<int> pos(n);
  for (int i = 0; i < n; i++) {
    pos[a[i]] = i; // a[i] -> pos[a[i]]
  }

  // The arrangement of [0, n-1] 
  //        is mapped to [1, n] in t
  ll ans = 0;
  FenWickTree<int> t(n + 1);
  for (int i = 0; i < n; i++) {
    int y = pos[b[i]];
    int lo = t.query(y);
    int hi = n - 1 - y - (i - lo);
    ans += (ll)lo * hi;
    t.add(y + 1, 1);
  }
  return ans;
}

复杂度

• 时间复杂度：\(O(nlogn)\)，优化了常数时间。
• 空间复杂度：\(O(n)\)