[LeetCode] 每日一题 2179. 统计数组中好三元组数目（新知识：树状数组）

题目描述

给你两个下标从 0 开始且长度为 n 的整数数组 nums1 和 nums2 ，两者都是 [0, 1, ..., n - 1] 的 排列 。

好三元组 指的是 3 个 互不相同 的值，且它们在数组 nums1 和 nums2 中出现顺序保持一致。换句话说，如果我们将 pos1v 记为值 v 在 nums1 中出现的位置，pos2v 为值 v 在 nums2 中的位置，那么一个好三元组定义为 0 <= x, y, z <= n - 1 ，且 pos1x < pos1y < pos1z 和 pos2x < pos2y < pos2z 都成立的 (x, y, z) 。

请你返回好三元组的 总数目 。

题目链接

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

示例输入

示例 1

输入：nums1 = [2,0,1,3], nums2 = [0,1,2,3]
输出：1
解释：
总共有 4 个三元组 (x,y,z) 满足 pos1x < pos1y < pos1z ，分别是 (2,0,1) ，(2,0,3) ，(2,1,3) 和 (0,1,3) 。
这些三元组中，只有 (0,1,3) 满足 pos2x < pos2y < pos2z 。所以只有 1 个好三元组。

示例 2

输入：nums1 = [4,0,1,3,2], nums2 = [4,1,0,2,3]
输出：4
解释：总共有 4 个好三元组 (4,0,3) ，(4,0,2) ，(4,1,3) 和 (4,1,2) 。

提示

• n == nums1.length == nums2.length

• 3 <= n <= 105

• 0 <= nums1[i], nums2[i] <= n - 1

• nums1 和 nums2 是 [0, 1, ..., n - 1] 的排列。

解题思路

这题相比昨天的“好三元组”难度明显上升了一些。两个数组 nums1 和 nums2 都是 0 ~ n-1 的排列，我们需要统计三元组 (x, y, z) 满足它们在两个数组中相对顺序一致。

从题意可以看出，这其实是要找两个数组的公共长度为3的上升子序列的个数。我一开始下意识就想用三维 dp 做法尝试解决，但不幸超时。后来参考了官方题解，发现可以用一种我之前不熟悉的数据结构 —— 树状数组（Fenwick Tree） 来优化求解过程。

核心想法是这样：

我们定义 indexMapping[i] = pos2[nums1[i]]，表示 nums1[i] 在 nums2 中的位置。这样问题就变成了：在数组 indexMapping 中统计有多少个三元组 (i, j, k) 满足 i < j < k 且 indexMapping[i] < indexMapping[j] < indexMapping[k]，这就是典型的求三元组递增子序列个数问题。

我们可以固定中间位置 j，对于每个 j，我们统计它左边比它小的个数 left，右边比它大的个数 right，它作为中间元素的三元组个数就是 left * right。

关键在于如何高效统计 left 和 right，这时就用到了树状数组：

• 我们按值从小到大遍历 indexMapping，在遍历中实时用树状数组维护前缀和，记录左边有多少个值比当前位置小；

• 右边比当前位置大的个数可以用 n - 1 - pos - (value - left) 推导得到；

• 最终把所有 left * right 累加即可。

代码实现

class FenwickTree {
    private int[] tree;

    public FenwickTree(int size) {
        tree = new int[size + 1];
    }

    public void update(int index, int delta) {
        index++;
        while (index < tree.length) {
            tree[index] += delta;
            index += index & -index;
        }
    }

    public int query(int index) {
        index++;
        int res = 0;
        while (index > 0) {
            res += tree[index];
            index -= index & -index;
        }
        return res;
    }
}

class Solution {
    public long goodTriplets(int[] nums1, int[] nums2) {
        int n = nums1.length;
        int[] pos2 = new int[n], reversedIndexMapping = new int[n];
        for (int i = 0; i < n; i++) {
            pos2[nums2[i]] = i;
        }
        for (int i = 0; i < n; i++) {
            reversedIndexMapping[pos2[nums1[i]]] = i;
        }
        FenwickTree tree = new FenwickTree(n);
        long res = 0;
        for (int value = 0; value < n; value++) {
            int pos = reversedIndexMapping[value];
            int left = tree.query(pos);
            tree.update(pos, 1);
            int right = (n - 1 - pos) - (value - left);
            res += (long) left * right;
        }
        return res;
    }
}

复杂度分析

• 时间复杂度：O(n log n)
  遍历 n 次，每次树状数组查询和更新都是 log 级别，整体非常高效 ⚡

• 空间复杂度：O(n)
  需要数组记录位置映射和构造树状数组 🌲

总结

这道题对我来说最大的收获是接触到了树状数组这种数据结构。在日常 Java 后端开发中不太容易遇到这种结构，但刷题时能认识这些工具是非常有价值的。树状数组本质上是一种能支持前缀和动态维护的数据结构，用于高效处理频繁查询和更新的问题。

如果你和我一样以前对树状数组不熟悉，这题就是一个很棒的入门例子 💡

希望这篇分享能为你带来启发！如果你有任何问题或建议，欢迎在评论区留言，与我共同交流探讨。