2179：统计数组中好三元组数目（2272 分）

力扣第 72 场双周赛第 4 题

题目

给你两个下标从 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) 。

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

示例 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 中的坐标，得到数组 A，问题转为求 A 中的递增三元组。

遍历每个中间数 y，找前面比 y 小的个数和后面比 y 大的个数即可求得 y 对应的三元组个数。

容易想到维护有序集合，二分查找即可。

解答

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

def goodTriplets(self, nums1: List[int], nums2: List[int]) -> int:
    from sortedcontainers import SortedList
    d = {num: i for i, num in enumerate(nums2)}
    A = [d[num] for num in nums1]
    res, sl1, sl2 = 0, SortedList(), SortedList(A)
    for y in A:
        res += sl1.bisect_left(y)*(len(sl2)-sl2.bisect_right(y))
        sl1.add(y)
        sl2.remove(y)
    return res

1728 ms