leetcode 3583. 统计特殊三元组 中等

给你一个整数数组 nums。

特殊三元组 定义为满足以下条件的下标三元组 (i, j, k)：

• 0 <= i < j < k < n，其中 n = nums.length
• nums[i] == nums[j] * 2
• nums[k] == nums[j] * 2

返回数组中 特殊三元组的总数。

由于答案可能非常大，请返回结果对 109 + 7 取余数后的值。

示例 1：

输入： nums = 6,3,6

输出： 1

解释：

唯一的特殊三元组是 (i, j, k) = (0, 1, 2)，其中：

• nums[0] = 6, nums[1] = 3, nums[2] = 6
• nums[0] = nums[1] * 2 = 3 * 2 = 6
• nums[2] = nums[1] * 2 = 3 * 2 = 6

示例 2：

输入： nums = 0,1,0,0

输出： 1

解释：

唯一的特殊三元组是 (i, j, k) = (0, 2, 3)，其中：

• nums[0] = 0, nums[2] = 0, nums[3] = 0
• nums[0] = nums[2] * 2 = 0 * 2 = 0
• nums[3] = nums[2] * 2 = 0 * 2 = 0

示例 3：

输入： nums = 8,4,2,8,4

输出： 2

解释：

共有两个特殊三元组：

• (i, j, k) = (0, 1, 3)
  • nums[0] = 8, nums[1] = 4, nums[3] = 8
  • nums[0] = nums[1] * 2 = 4 * 2 = 8
  • nums[3] = nums[1] * 2 = 4 * 2 = 8
• (i, j, k) = (1, 2, 4)
  • nums[1] = 4, nums[2] = 2, nums[4] = 4
  • nums[1] = nums[2] * 2 = 2 * 2 = 4
  • nums[4] = nums[2] * 2 = 2 * 2 = 4

提示：

• 3 <= n == nums.length <= 10^5
• 0 <= nums[i] <= 10^5

分析：枚举中间的数 numj，检查左边和右边各有多少个等于 numj*2 的数，记为 left 和 right，则位置 j 的 numj 能组成的三元组个数为 left*right。

class Solution {
public:
    int specialTriplets(vector<int>& nums) {
        vector<vector<int>>ind(100005);
        int ans=0,mod=1e9+7,n=nums.size(),cnt=1;
        map<int,int>mp;
        for(int i=0;i<n;++i)
        {
            if(mp.find(nums[i])==mp.end())mp[nums[i]]=cnt++;
            ind[mp[nums[i]]].push_back(i);
        }

        for(int i=1;i<n;++i)
        {
            if(mp.find(nums[i]*2)!=mp.end())
            {
                int left=lower_bound(ind[mp[nums[i]*2]].begin(),ind[mp[nums[i]*2]].end(),i)-ind[mp[nums[i]*2]].begin();
                int right=upper_bound(ind[mp[nums[i]*2]].begin(),ind[mp[nums[i]*2]].end(),i)-ind[mp[nums[i]*2]].begin();
                ans=(left*1LL*(ind[mp[nums[i]*2]].size()-right)%mod+ans)%mod;
            }
        }
        return ans;
    }
};

上面的代码有些想复杂了，可以用一个 map 记录所有数字的出现次数，每次枚举到位置 j 时，把这个位置的计数从 map 里去掉，再用一个 map，记录当前位置之前所有数字出现的次数，这样就不需要用 vector 记录位置。

class Solution {
public:
    int specialTriplets(vector<int>& nums) {
        int ans=0,mod=1e9+7,n=nums.size(),cnt=1;
        map<int,int>pre,last;
        for(int i=0;i<n;++i)
            last[nums[i]]++;

        last[nums[0]]--;pre[nums[0]]++;
        for(int i=1;i<n-1;++i)
        {
            last[nums[i]]--;
            ans=(ans+pre[nums[i]*2]*1LL*last[nums[i]*2]%mod)%mod;
            pre[nums[i]]++;
        }
        return ans;
    }
};