LeetCode-in-Java

3583. Count Special Triplets

Medium

You are given an integer array nums.

A special triplet is defined as a triplet of indices (i, j, k) such that:

Return the total number of special triplets in the array.

Since the answer may be large, return it modulo 109 + 7.

Example 1:

Input: nums = [6,3,6]

Output: 1

Explanation:

The only special triplet is (i, j, k) = (0, 1, 2), where:

Example 2:

Input: nums = [0,1,0,0]

Output: 1

Explanation:

The only special triplet is (i, j, k) = (0, 2, 3), where:

Example 3:

Input: nums = [8,4,2,8,4]

Output: 2

Explanation:

There are exactly two special triplets:

Constraints:

Solution

public class Solution {
    public int specialTriplets(int[] nums) {
        long ans = 0;
        int[] sum = new int[200002];
        int[] left = new int[nums.length];
        for (int i = 0; i < nums.length; i++) {
            int curr = nums[i];
            sum[curr]++;
            left[i] = sum[curr * 2];
        }
        for (int i = 0; i < nums.length; i++) {
            int curr = nums[i];
            long leftCount = curr == 0 ? left[i] - 1 : left[i];
            long rightCount = sum[curr * 2] - (long) left[i];
            if (leftCount != 0 && rightCount != 0) {
                ans += leftCount * rightCount;
            }
        }
        return (int) (ans % 1000000007);
    }
}