Count Good Triplets in an Array

Hard

#1984Time: O(n log n). The pre-computation takes O(n). The main loop runs `n` times, and inside the loop, the Fenwick Tree operations (query and update) each take O(log n) time. This results in a total time complexity of O(n log n).Space: O(n) for the `pos2` map, the transformed array `A`, and the Fenwick Tree.1 company

Algorithms

Binary Search Divide and Conquer Merge Sort

Data structures

Array Binary Indexed Tree Segment Tree Ordered Set

Companies

Walmart Labs

Prompt

You are given two 0-indexed arrays nums1 and nums2 of length n, both of which are permutations of [0, 1, ..., n - 1].

A good triplet is a set of 3 distinct values which are present in increasing order by position both in nums1 and nums2. In other words, if we consider pos1_v as the index of the value v in nums1 and pos2_v as the index of the value v in nums2, then a good triplet will be a set (x, y, z) where 0 <= x, y, z <= n - 1, such that pos1_x < pos1_y < pos1_z and pos2_x < pos2_y < pos2_z.

Return the total number of good triplets.

Example 1:

Input: nums1 = [2,0,1,3], nums2 = [0,1,2,3]
Output: 1
Explanation: 
There are 4 triplets (x,y,z) such that pos1x < pos1y < pos1z. They are (2,0,1), (2,0,3), (2,1,3), and (0,1,3). 
Out of those triplets, only the triplet (0,1,3) satisfies pos2x < pos2y < pos2z. Hence, there is only 1 good triplet.

Example 2:

Input: nums1 = [4,0,1,3,2], nums2 = [4,1,0,2,3]
Output: 4
Explanation: The 4 good triplets are (4,0,3), (4,0,2), (4,1,3), and (4,1,2).

Constraints:

n == nums1.length == nums2.length
3 <= n <= 10⁵
0 <= nums1[i], nums2[i] <= n - 1
nums1 and nums2 are permutations of [0, 1, ..., n - 1].

Approaches

2 approaches with complexity analysis and trade-offs.

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.

Algorithm

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).

Walkthrough

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.

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.
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).
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.
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.

1class Solution {2    // Fenwick Tree (Binary Indexed Tree) implementation3    class FenwickTree {4        private int[] bit;5        private int size;6 7        public FenwickTree(int size) {8            this.size = size;9            this.bit = new int[size + 1];10        }11 12        public void update(int index, int delta) {13            while (index <= size) {14                bit[index] += delta;15                index += index & -index; // Move to the next relevant index16            }17        }18 19        public int query(int index) {20            int sum = 0;21            while (index > 0) {22                sum += bit[index];23                index -= index & -index; // Move to the parent index24            }25            return sum;26        }27    }28 29    public long goodTriplets(int[] nums1, int[] nums2) {30        int n = nums1.length;31        int[] pos2 = new int[n];32        for (int i = 0; i < n; ++i) {33            pos2[nums2[i]] = i;34        }35 36        int[] A = new int[n];37        for (int i = 0; i < n; ++i) {38            A[i] = pos2[nums1[i]];39        }40 41        // The values in A are 0 to n-1. BIT needs size n+1 for 1-based indexing.42        FenwickTree bit = new FenwickTree(n);43        long totalTriplets = 0;44 45        for (int j = 0; j < n; ++j) {46            // A[j] is the middle element. Values are 0-indexed.47            // Query for elements < A[j]. BIT is 1-indexed, so query up to A[j].48            int smaller_left = bit.query(A[j]);49            50            // Number of elements processed so far (to the left of j) is j.51            // Number of elements > A[j] on the left.52            int larger_left = j - smaller_left;53 54            // Total numbers in the permutation > A[j] is (n - 1 - A[j]).55            // Number of elements > A[j] on the right.56            int larger_right = (n - 1 - A[j]) - larger_left;57 58            totalTriplets += (long) smaller_left * larger_right;59 60            // Add current element A[j] to BIT for subsequent calculations.61            // Update at index A[j] + 1 because BIT is 1-indexed.62            bit.update(A[j] + 1, 1);63        }64 65        return totalTriplets;66    }67}

Complexity

Time

O(n log n). The pre-computation takes O(n). The main loop runs `n` times, and inside the loop, the Fenwick Tree operations (query and update) each take O(log n) time. This results in a total time complexity of O(n log n).

Space

O(n) for the `pos2` map, the transformed array `A`, and the Fenwick Tree.

Trade-offs

Pros

Highly efficient with O(n log n) time complexity, which passes the given constraints.
It's a standard and powerful technique for solving a class of problems related to counting inversions or ordered subsequences.

Cons

Requires knowledge of advanced data structures like Fenwick Trees or Segment Trees.
Implementation is more complex than the brute-force approach.

This approach first simplifies the problem. The condition of having increasing positions in both arrays is complex. We can fix the order of one array, say nums1, and then check the condition on the second. By iterating through nums1, any three values x, y, z picked at indices i < j < k automatically satisfy pos1_x < pos1_y < pos1_z. The problem then reduces to counting how many of these triplets also satisfy pos2_x < pos2_y < pos2_z.

To make this check efficient, we transform nums1 into a new array A where A[i] stores the position of nums1[i] in nums2. Now, the problem is equivalent to finding the number of increasing subsequences of length 3 in A.

The brute-force method iterates through every possible middle element A[j] of a triplet, and for each, it counts the number of smaller elements to its left and the number of larger elements to its right by linear scans.

Algorithm

Create a mapping pos2 from each value to its index in nums2. This allows for O(1) lookup of a value's position in nums2.
Create a new array A of length n, where A[i] = pos2[nums1[i]]. This transformation simplifies the problem: a triplet of values (x, y, z) is good if and only if their corresponding indices i, j, k in A satisfy i < j < k and A[i] < A[j] < A[k].
The problem is now to count the number of increasing subsequences of length 3 in array A.
Initialize a long variable count to 0.
Iterate through the array A with an index j from 1 to n-2. This A[j] will be the middle element of a potential triplet.
For each j, we need to find:
- smaller_left: The number of elements A[i] with i < j such that A[i] < A[j].
- larger_right: The number of elements A[k] with k > j such that A[k] > A[j].
To find smaller_left, iterate with an index i from 0 to j-1 and count.
To find larger_right, iterate with an index k from j+1 to n-1 and count.
The number of good triplets with A[j] as the middle element is smaller_left * larger_right.
Add this product to the total count.
After the main loop finishes, return count.

Walkthrough

First, we need a way to quickly find the position of any number in nums2. We can pre-process nums2 and store the indices in a hash map or an array, let's call it pos2. pos2[value] = index.

Next, we create a new array A. For each index i from 0 to n-1, we set A[i] = pos2[nums1[i]]. Now, if we pick three indices i < j < k, the values x=nums1[i], y=nums1[j], z=nums1[k] satisfy the first condition (pos1_x < pos1_y < pos1_z). For these to form a good triplet, they must also satisfy pos2_x < pos2_y < pos2_z, which translates to pos2[nums1[i]] < pos2[nums1[j]] < pos2[nums1[k]], or simply A[i] < A[j] < A[k].

So, the task is to count triplets of indices (i, j, k) such that i < j < k and A[i] < A[j] < A[k]. We can do this by iterating through each possible middle index j from 1 to n-2. For each j, we count how many elements to its left are smaller (smaller_left) and how many elements to its right are larger (larger_right). The number of triplets centered at j is then smaller_left * larger_right. Summing this product over all possible j gives the total count.

1class Solution {2    public long goodTriplets(int[] nums1, int[] nums2) {3        int n = nums1.length;4        int[] pos2 = new int[n];5        for (int i = 0; i < n; ++i) {6            pos2[nums2[i]] = i;7        }8 9        int[] A = new int[n];10        for (int i = 0; i < n; ++i) {11            A[i] = pos2[nums1[i]];12        }13 14        long totalTriplets = 0;15        for (int j = 1; j < n - 1; ++j) {16            int smaller_left = 0;17            for (int i = 0; i < j; ++i) {18                if (A[i] < A[j]) {19                    smaller_left++;20                }21            }22 23            int larger_right = 0;24            for (int k = j + 1; k < n; ++k) {25                if (A[k] > A[j]) {26                    larger_right++;27                }28            }29            totalTriplets += (long) smaller_left * larger_right;30        }31 32        return totalTriplets;33    }34}

Complexity

Time

O(n^2). The pre-computation of `pos2` and `A` takes O(n). The main part involves a nested loop structure. The outer loop for `j` runs `n` times, and for each `j`, the inner loops for `i` and `k` take O(j) and O(n-j) time respectively, leading to an overall O(n^2) complexity.

Space

O(n) to store the `pos2` mapping and the transformed array `A`.

Trade-offs

Pros

Relatively simple to understand and implement.
Improves upon a naive O(n^3) solution by fixing the middle element.

Cons

The time complexity of O(n^2) is too slow for the given constraints where n can be up to 10^5, and will result in a 'Time Limit Exceeded' error.

Solutions

Solution

1class Solution {2public3  long goodTriplets(int[] nums1, int[] nums2) {4    int n = nums1.length;5    int[] pos = new int[n];6    BinaryIndexedTree tree = new BinaryIndexedTree(n);7    for (int i = 0; i < n; ++i) {8      pos[nums2[i]] = i + 1;9    }10    long ans = 0;11    for (int num : nums1) {12      int p = pos[num];13      long left = tree.query(p);14      long right = n - p - (tree.query(n) - tree.query(p));15      ans += left * right;16      tree.update(p, 1);17    }18    return ans;19  }20} class BinaryIndexedTree {21private22  int n;23private24  int[] c;25public26  BinaryIndexedTree(int n) {27    this.n = n;28    c = new int[n + 1];29  }30public31  void update(int x, int delta) {32    while (x <= n) {33      c[x] += delta;34      x += lowbit(x);35    }36  }37public38  int query(int x) {39    int s = 0;40    while (x > 0) {41      s += c[x];42      x -= lowbit(x);43    }44    return s;45  }46public47  static int lowbit(int x) { return x & -x; }48}

1class BinaryIndexedTree { public: int n ; vector < int > c ; BinaryIndexedTree ( int _n ) : n ( _n ) , c ( _n + 1 ) {} void update ( int x , int delta ) { while ( x <= n ) { c [ x ] += delta ; x += lowbit ( x ); } } int query ( int x ) { int s = 0 ; while ( x > 0 ) { s += c [ x ]; x -= lowbit ( x ); } return s ; } int lowbit ( int x ) { return x & - x ; } }; class Solution { public: long long goodTriplets ( vector < int >& nums1 , vector < int >& nums2 ) { int n = nums1 . size (); vector < int > pos ( n ); for ( int i = 0 ; i < n ; ++ i ) pos [ nums2 [ i ]] = i + 1 ; BinaryIndexedTree * tree = new BinaryIndexedTree ( n ); long long ans = 0 ; for ( int & num : nums1 ) { int p = pos [ num ]; int left = tree -> query ( p ); int right = n - p - ( tree -> query ( n ) - tree -> query ( p )); ans += 1ll * left * right ; tree -> update ( p , 1 ); } return ans ; } };

1class BinaryIndexedTree : def __init__ ( self , n ): self . n = n self . c = [ 0 ] * ( n + 1 ) @ staticmethod def lowbit ( x ): return x & - x def update ( self , x , delta ): while x <= self . n : self . c [ x ] += delta x += BinaryIndexedTree . lowbit ( x ) def query ( self , x ): s = 0 while x > 0 : s += self . c [ x ] x -= BinaryIndexedTree . lowbit ( x ) return s class Solution : def goodTriplets ( self , nums1 : List [ int ], nums2 : List [ int ]) -> int : pos = { v : i for i , v in enumerate ( nums2 , 1 )} ans = 0 n = len ( nums1 ) tree = BinaryIndexedTree ( n ) for num in nums1 : p = pos [ num ] left = tree . query ( p ) right = n - p - ( tree . query ( n ) - tree . query ( p )) ans += left * right tree . update ( p , 1 ) return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem