Count Good Triplets in an Array

The O(n^2) approach is too slow because counting smaller elements to the left and larger elements to the right takes O(n) time for each middle element. This counting process can be optimized using a data structure that supports fast queries and updates.

A Fenwick Tree (or Binary Indexed Tree, BIT) is ideal for this. It can count the number of elements in a range and add a new element in O(log n) time. By processing the elements of the transformed array A and using a BIT, we can calculate the required counts efficiently.

We can either use two passes over the array A—one forward and one backward, each with its own BIT—to compute arrays of smaller_left and larger_right counts. Or, more elegantly, we can devise a single-pass algorithm that calculates the contribution of each element as it iterates, using just one BIT.

• First, perform the same problem transformation as in the brute-force approach. Create pos2 map and the array A where A[i] = pos2[nums1[i]].
• The goal is to efficiently count, for each j, the number of smaller elements to the left (smaller_left) and larger elements to the right (larger_right).
• Two-Pass Method:
  1. Calculate smaller_left for all j: Initialize a Fenwick Tree (BIT). Iterate j from 0 to n-1. For each j, smaller_left[j] is the count of numbers less than A[j] seen so far, which can be found with BIT.query(A[j]). Then, update the BIT with A[j] by calling BIT.update(A[j] + 1, 1).
  2. Calculate larger_right for all j: Initialize a new BIT. Iterate j from n-1 down to 0. For each j, larger_right[j] is the count of numbers greater than A[j] seen so far (i.e., to the right). This is (elements processed so far) - (elements <= A[j]), which can be found with BIT.query(n) - BIT.query(A[j] + 1). Then, update the BIT with A[j].
  3. Finally, iterate from j=0 to n-1 and sum up the products (long)smaller_left[j] * larger_right[j].
• Single-Pass Optimization:
  1. We can compute the total count in a single pass. Initialize one BIT and a total_triplets counter.
  2. Iterate j from 0 to n-1.
  3. In each iteration, A[j] is the middle element. The elements A[0...j-1] are already processed and stored in the BIT.
  4. smaller_left = BIT.query(A[j]). This is the count of elements seen so far that are smaller than A[j].
  5. larger_right can be calculated without a second pass. The total count of numbers in the universe [0...n-1] that are greater than A[j] is (n - 1 - A[j]). Some of these are on the left of j, and some are on the right. The count of numbers greater than A[j] on the left is larger_left = (j) - smaller_left. So, larger_right = (total_greater) - (larger_left) = (n - 1 - A[j]) - (j - smaller_left).
  6. Add (long)smaller_left * larger_right to total_triplets.
  7. Update the BIT with A[j] for future iterations: BIT.update(A[j] + 1, 1).

The core idea is to optimize the counting of smaller_left and larger_right for each potential middle element A[j]. After transforming the problem to finding increasing subsequences of length 3 in array A, we can proceed with a single-pass approach using a Fenwick Tree.

We iterate through A from left to right (index j). For each element A[j], we treat it as the middle element of a triplet (A[i], A[j], A[k]) where i < j < k.

1. Count smaller elements on the left (smaller_left): The elements A[0], ..., A[j-1] have already been processed. We can use a Fenwick Tree to maintain the counts of these numbers. A query to the BIT, BIT.query(A[j]), will give us the number of elements seen so far that are strictly less than A[j]. This is our smaller_left count.

2. Count larger elements on the right (larger_right): We need to count elements A[k] where k > j and A[k] > A[j]. We can deduce this count without iterating. The total number of values in the permutation greater than A[j] is (n - 1) - A[j]. These values are distributed to the left and right of index j. We can find how many are on the left: larger_left is the total number of elements to the left (j) minus those that are smaller (smaller_left). So, larger_left = j - smaller_left. Therefore, larger_right = (total_values > A[j]) - larger_left = (n - 1 - A[j]) - (j - smaller_left).

3. Update total: For the current middle element A[j], we have found smaller_left and larger_right. We add their product to our running total.

4. Update BIT: After processing A[j], we must add it to our Fenwick Tree so it's included in the counts for subsequent elements. We do this with BIT.update(A[j] + 1, 1).

This single pass correctly accumulates the total count of good triplets.

class Solution {
    // Fenwick Tree (Binary Indexed Tree) implementation
    class FenwickTree {
        private int[] bit;
        private int size;

        public FenwickTree(int size) {
            this.size = size;
            this.bit = new int[size + 1];
        }

        public void update(int index, int delta) {
            while (index <= size) {
                bit[index] += delta;
                index += index & -index; // Move to the next relevant index
            }
        }

        public int query(int index) {
            int sum = 0;
            while (index > 0) {
                sum += bit[index];
                index -= index & -index; // Move to the parent index
            }
            return sum;
        }
    }

    public long goodTriplets(int[] nums1, int[] nums2) {
        int n = nums1.length;
        int[] pos2 = new int[n];
        for (int i = 0; i < n; ++i) {
            pos2[nums2[i]] = i;
        }

        int[] A = new int[n];
        for (int i = 0; i < n; ++i) {
            A[i] = pos2[nums1[i]];
        }

        // The values in A are 0 to n-1. BIT needs size n+1 for 1-based indexing.
        FenwickTree bit = new FenwickTree(n);
        long totalTriplets = 0;

        for (int j = 0; j < n; ++j) {
            // A[j] is the middle element. Values are 0-indexed.
            // Query for elements < A[j]. BIT is 1-indexed, so query up to A[j].
            int smaller_left = bit.query(A[j]);
            
            // Number of elements processed so far (to the left of j) is j.
            // Number of elements > A[j] on the left.
            int larger_left = j - smaller_left;

            // Total numbers in the permutation > A[j] is (n - 1 - A[j]).
            // Number of elements > A[j] on the right.
            int larger_right = (n - 1 - A[j]) - larger_left;

            totalTriplets += (long) smaller_left * larger_right;

            // Add current element A[j] to BIT for subsequent calculations.
            // Update at index A[j] + 1 because BIT is 1-indexed.
            bit.update(A[j] + 1, 1);
        }

        return totalTriplets;
    }
}