Hard

题目描述

给定两个长度为 n 的下标从 0 开始的数组 nums1 和 nums2,它们都是 [0, 1, …, n - 1] 的排列。

好三元组是一组 3 个不同的值,这些值在 nums1 和 nums2 中都按照位置递增的顺序出现。换句话说,如果我们将 pos1_v 视为值 v 在 nums1 中的下标,将 pos2_v 视为值 v 在 nums2 中的下标,那么好三元组就是一个集合 (x, y, z),其中 0 <= x, y, z <= n - 1,满足 pos1_x < pos1_y < pos1_z 且 pos2_x < pos2_y < pos2_z。

返回好三元组的总数。

示例 1:

输入:nums1 = [2,0,1,3], nums2 = [0,1,2,3]
输出:1
解释:
存在 4 个三元组 (x,y,z) 满足 pos1_x < pos1_y < pos1_z。它们是 (2,0,1), (2,0,3), (2,1,3) 和 (0,1,3)。
在这些三元组中,只有三元组 (0,1,3) 满足 pos2_x < pos2_y < pos2_z。因此,只有 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 <= 10^5
  • 0 <= nums1[i], nums2[i] <= n - 1
  • nums1 和 nums2 是 [0, 1, …, n - 1] 的排列

提示:

  • 对于每个值 y,如何找到在两个数组中都出现在 y 之前的值 x 的数量?
  • 类似地,对于每个值 y,找到在两个数组中都出现在 y 之后的值 z 的数量。
  • 现在,对于每个值 y,计算以 y 作为中间元素可以形成的好三元组数量。

解题思路

这道题的核心思路是固定中间元素,然后计算能够形成好三元组的数量。

主要思路:

  1. 首先记录每个值在两个数组中的位置
  2. 对于每个可能的中间元素 y,我们需要找到:
    • 在两个数组中都出现在 y 之前的元素个数(称为 left)
    • 在两个数组中都出现在 y 之后的元素个数(称为 right)
  3. 以 y 为中间元素的好三元组数量就是 left × right

实现细节:

  • 可以使用树状数组(Binary Indexed Tree)或者归并排序来高效计算逆序对
  • 另一种方法是对于每个位置,计算在两个数组中都满足条件的前缀和后缀数量

优化: 本题使用树状数组是比较高效的方法,时间复杂度为 O(n log n)。我们遍历 nums1,对于每个位置的元素,计算它在 nums2 中对应位置之前有多少个已经访问过的元素(这些就是满足条件的左边元素),以及之后有多少个未访问的元素(这些是满足条件的右边元素)。

代码实现

class Solution {
private:
    vector<int> tree;
    int n;
    
    void update(int idx, int val) {
        for (int i = idx + 1; i <= n; i += i & -i) {
            tree[i] += val;
        }
    }
    
    int query(int idx) {
        int sum = 0;
        for (int i = idx + 1; i > 0; i -= i & -i) {
            sum += tree[i];
        }
        return sum;
    }
    
public:
    long long goodTriplets(vector<int>& nums1, vector<int>& nums2) {
        n = nums1.size();
        tree.resize(n + 1, 0);
        
        vector<int> pos1(n), pos2(n);
        for (int i = 0; i < n; i++) {
            pos1[nums1[i]] = i;
            pos2[nums2[i]] = i;
        }
        
        long long result = 0;
        
        for (int i = 0; i < n; i++) {
            int val = nums1[i];
            int pos_in_nums2 = pos2[val];
            
            // 计算在nums2中位置小于pos_in_nums2且已经在nums1中访问过的元素个数
            long long left = query(pos_in_nums2 - 1);
            
            // 计算在nums2中位置大于pos_in_nums2且还未在nums1中访问过的元素个数
            long long right = (n - 1 - pos_in_nums2) - (i - left);
            
            result += left * right;
            
            // 更新树状数组
            update(pos_in_nums2, 1);
        }
        
        return result;
    }
};
class Solution:
    def goodTriplets(self, nums1: List[int], nums2: List[int]) -> int:
        n = len(nums1)
        
        # 记录每个值在两个数组中的位置
        pos1 = [0] * n
        pos2 = [0] * n
        for i in range(n):
            pos1[nums1[i]] = i
            pos2[nums2[i]] = i
        
        # 树状数组
        class BIT:
            def __init__(self, size):
                self.size = size
                self.tree = [0] * (size + 1)
            
            def update(self, idx, val):
                idx += 1
                while idx <= self.size:
                    self.tree[idx] += val
                    idx += idx & -idx
            
            def query(self, idx):
                idx += 1
                result = 0
                while idx > 0:
                    result += self.tree[idx]
                    idx -= idx & -idx
                return result
        
        bit = BIT(n)
        result = 0
        
        for i in range(n):
            val = nums1[i]
            pos_in_nums2 = pos2[val]
            
            # 在nums2中位置小于当前位置且已访问的元素个数
            left = bit.query(pos_in_nums2 - 1) if pos_in_nums2 > 0 else 0
            
            # 在nums2中位置大于当前位置且未访问的元素个数
            right = (n - 1 - pos_in_nums2) - (i - left)
            
            result += left * right
            bit.update(pos_in_nums2, 1)
        
        return result
public class Solution {
    private int[] tree;
    private int n;
    
    private void Update(int idx, int val) {
        for (int i = idx + 1; i <= n; i += i & -i) {
            tree[i] += val;
        }
    }
    
    private int Query(int idx) {
        int sum = 0;
        for (int i = idx + 1; i > 0; i -= i & -i) {
            sum += tree[i];
        }
        return sum;
    }
    
    public long GoodTriplets(int[] nums1, int[] nums2) {
        n = nums1.Length;
        tree = new int[n + 1];
        
        int[] pos1 = new int[n];
        int[] pos2 = new int[n];
        
        for (int i = 0; i < n; i++) {
            pos1[nums1[i]] = i;
            pos2[nums2[i]] = i;
        }
        
        long result = 0;
        
        for (int i = 0; i < n; i++) {
            int val = nums1[i];
            int posInNums2 = pos2[val];
            
            long left = posInNums2 > 0 ? Query(posInNums2 - 1) : 0;
            long right = (n - 1 - posInNums2) - (i - left);
            
            result += left * right;
            Update(posInNums2, 1);
        }
        
        return result;
    }
}
var goodTriplets = function(nums1, nums2) {
    const n = nums1.length;
    
    // 记录每个值在两个数组中的位置
    const pos1 = new Array(n);
    const pos2 = new Array(n);
    for (let i = 0; i < n; i++) {
        pos1[nums1[i]] = i;
        pos2[nums2[i]] = i;
    }
    
    // 树状数组
    class BIT {
        constructor(size) {
            this.size = size;
            this.tree = new Array(size + 1).fill(0);
        }
        
        update(idx, val) {
            idx++;
            while (idx <= this.size) {
                this.tree[idx] += val;
                idx += idx & -idx;
            }
        }
        
        query(idx) {
            idx++;
            let result = 0;
            while (idx > 0) {
                result += this.tree[idx];
                idx -= idx & -idx;
            }
            return result;
        }
    }
    
    const bit = new BIT(n);
    let result = 0;
    
    for (let i = 0; i < n; i++) {
        const val = nums1[i];
        const posInNums2 = pos2[val];
        
        const left = posInNums2 > 0 ? bit.query(posInNums2 - 1) : 0;
        const right = (n - 1 - posInNums2) - (i - left);
        
        result += left * right;
        bit.update(posInNums2, 1);
    }
    
    return result;
};

复杂度分析

复杂度类型复杂度说明
时间复杂度O(n log n)需要进行 n 次树状数组的更新和查询操作,每次操作时间复杂度为 O(log n)
空间复杂度O(n)需要额外的数组存储位置信息和树状数组

相关题目