2179. Count Good Triplets in an Array

HardBinary Indexed TreeSegment TreeArrayBinary SearchDivide and ConquerOrdered SetMerge Sort

Problem Description

The problem requires us to find "good triplets" in two given arrays, nums1 and nums2, both of which are permutations of [0, 1, ..., n - 1]. A "good triplet" is defined as a set of three distinct values (x, y, z) which are in increasing order by their indices in both nums1 and nums2. To be more specific, we look for values such that their indices pos1_x, pos1_y, and pos1_z in nums1 and pos2_x, pos2_y, and pos2_z in nums2 follow these conditions: pos1_x < pos1_y < pos1_z and pos2_x < pos2_y < pos2_z. The problem asks us to return the total number of these good triplets.

Intuition

The key to solving this problem is recognizing that brute force checking all possible triplets would be inefficient, as it would require O(n^3) time complexity. Instead, we can use a Binary Indexed Tree (BIT) or Segment Tree, data structures that allow us to efficiently update elements and query the sum or minimum/maximum of elements in a range. In the context of this problem, we will use a Binary Indexed Tree to maintain the count of numbers while iterating through nums1.

The intuition behind using a BIT is as follows:

We want to determine, for each element in nums1, how many elements before it would form a good pair (the x and y of a good triplet where this element is z).
For each element, we also need to find out how many numbers would come after it in nums2 to serve as potential z's for the good pair identified in step 1.
To find the number of elements that precede a given element in nums2, we can use the BIT to perform prefix sums (i.e., count the number of elements we've seen so far in nums2 that are less than our current element's position). The prefix sum essentially tells us how many y's there could be for a given x where x is an element we've seen earlier.
Then, for the elements that will come after our current element, we take the total number of elements (n) and subtract both the position of our current element and the number of elements that precede it (counted in step 3). This step gives us the number of possible z's for our current element being a y.
For every number in nums1, we calculate the product of the number of possible y's and z's as found in steps 3 and 4 to get the count of good triplets for which this is the middle element (y).
We sum these counts to get the total count of good triplets and update the BIT as we iterate through nums1, adding each element as we go.

With this approach, we avoid checking every possible triplet and work with an efficient update and query operation, resulting in a solution that is significantly faster than the brute force approach and suitable for handling permutations of arrays with large numbers of elements.

Learn more about Segment Tree, Binary Search, Divide and Conquer and Merge Sort patterns.

Solution Approach

The solution leverages a data structure called a Binary Indexed Tree (BIT), also known as a Fenwick Tree, to efficiently compute the number of good triplets as it iterates through nums1. Here's a detailed explanation of how the given Python class BinaryIndexedTree and the Solution class's goodTriplets method work together:

Binary Indexed Tree (BIT): This data structure supports efficient updates as well as prefix sum queries in logarithmic time. It's particularly useful when you need to perform frequent updates and retrieve aggregate information over a sequence of values.
The BinaryIndexedTree Class:
- The constructor initializes an array c which will contain the tree representation with all elements initialized to zero.
- lowbit is a helper static method that calculates the lowest bit of x which is essentially x ANDed with its two's complement (-x). This returns the value of the least significant bit that is set to one.
- The update method updates the tree by adding delta to position x. It starts at x and increments the index by lowbit(x), which moves to the subsequent value that needs to be updated due to the change at x.
- The query method calculates a prefix sum up to index x. It uses lowbit to move through the indices that contribute to the prefix sum up to x.
The goodTriplets Method in Solution Class:
- A dictionary pos is made to record the position of each value in nums2, shifted by one for BIT indexing purposes (since BIT is typically 1-indexed).
- An instance of BinaryIndexedTree is created with size n, which is the length of the given permutations.
- The main loop iterates over each number in nums1 to compute the number of good triplets it contributes to.
- For a current number, query the BIT using its position in nums2 to get left, the number of elements in nums1 that have appeared before the current element in nums2.
- Calculate right by subtracting the number of elements that came before and including the current element from the total number of elements, giving the number of elements that can be to the right of the current element in a good triplet.
- The multiplication of left and right gives the number of good triplets that can be formed where the current element is the middle element (y).
- Increment ans by this product.
- Finally, update the BIT to include this current element using the update method.
- After all elements are processed, return ans which contains the total count of good triplets.

By using a BIT, this approach significantly reduces the complexity of updating counts and computing prefix sums, allowing the computation of good triplets efficiently as the solution iterates through nums1. The overall complexity of this solution is O(n log n) due to each BIT operation taking O(log n) time and being performed n times.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's say we have the following two arrays, nums1 and nums2, both permutations of [0, 1, 2, 3, 4]:

nums1 = [4, 0, 3, 1, 2]
nums2 = [1, 3, 0, 2, 4]

We are tasked to find "good triplets" according to their description in the problem description.

Firstly, we create a Dictionary pos from nums2 to find the positions of each number quickly. For nums2 above, pos would be:

pos = {1: 1, 3: 2, 0: 3, 2: 4, 4: 5}   # Shifted by 1 for 1-indexing in BIT

Now we need to iterate through each element in nums1 and use the BIT to compute our results.

We have initialized an instance of BinaryIndexedTree called bit with length n+1 (since our arrays are zero-indexed, but the BIT is one-indexed).

Let's begin with processing nums1:

We look at the first number, 4. For 4, the position in nums2 according to pos is 5. There are zero elements before it in nums2 that appear in nums1, so left is 0 (no elements in nums1 appeared before 4 in nums2). There will be no elements after 4 in nums2 because it’s the last element. Thus, the number of good triplets with 4 as the middle element is 0 * (total - 5 - 0) = 0.
Now we update bit for 4 with bits.update(5, 1) so future queries will consider 4.

Continue for next number in nums1, 0:

pos[0] is 3. Query bit to get the count of elements in nums1 before 0 appeared in nums2, which is bit.query(3) = 0. The number of elements that can come after 0 is 5 - 3 - 0 = 2. So, good triplets with 0 as the middle are 0 * 2 = 0.

Update bit for 0 with bits.update(3, 1).

Continuing this process for the following numbers in nums1:

For 3: We query bit.query(2) = 1 (as 1 has appeared before 3 in nums2); the possible z values are 5 - 2 - 1 = 2. The good triplets with 3 as the middle are 1 * 2 = 2.
Update bit for 3 with bits.update(2, 1).
For 1: We query bit.query(1) = 0 (no elements appeared before 1 in nums2); the possible z values are 5 - 1 - 0 = 4. The good triplets with 1 as the middle are 0 * 4 = 0. Update bit for 1 with bits.update(1, 1).
For 2: We query bit.query(4) = 2 (as 1 and 3 have appeared before 2 in nums2); the possible z values are 5 - 4 - 2 = -1. Since this is not possible, we conclude there are no good triplets with 2 as the middle.

By the end, we sum the good triplets formed with each number in nums1 as the middle element, which totals to 2 in this example. Thus, the solution returns 2.

This walkthrough with a small example illustrates the efficient use of Binary Indexed Tree data structure in finding "good triplets" in two sequences.

Solution Implementation

1class BinaryIndexedTree:
2    # Initialization of the binary indexed tree with given size n.
3    def __init__(self, size):
4        self.size = size
5        self.tree_array = [0] * (size + 1)
6
7    # Calculate the least significant bit that is set (i.e., lowest power of 2 that divides x).
8    @staticmethod
9    def lowbit(x):
10        return x & -x
11
12    # Update the binary indexed tree by adding `delta` to index `index`.
13    def update(self, index, delta):
14        while index <= self.size:
15            self.tree_array[index] += delta
16            index += BinaryIndexedTree.lowbit(index)
17
18    # Compute the prefix sum from the start to index `index`.
19    def query(self, index):
20        sum_value = 0
21        while index > 0:
22            sum_value += self.tree_array[index]
23            index -= BinaryIndexedTree.lowbit(index)
24        return sum_value
25
26
27class Solution:
28    # Function to count the number of "good triplets" in the permutation of nums1 that satisfy the condition
29    # in the problem statement, using nums2 to determine the ideal order.
30    def goodTriplets(self, nums1: List[int], nums2: List[int]) -> int:
31        # Build a position mapping of elements in nums2 for O(1) lookup.
32        position = {value: index for index, value in enumerate(nums2, 1)}
33        good_triplets_count = 0  # Initialize counter for good triplets
34        length = len(nums1)  # The length of nums1 and nums2 arrays
35        tree = BinaryIndexedTree(length)  # Instantiate a Binary Indexed Tree
36
37        # Process each number in nums1 to count the number of good triplets.
38        for number in nums1:
39            order = position[number]  # Find the position of the current number in nums2
40            left_count = tree.query(order)  # Count of numbers before the current number's ideal position
41            # Count of numbers after the current number by subtracting the counts obtained from the tree
42            right_count = length - order - (tree.query(length) - tree.query(order))
43            # Increment the count of good triplets by the product of left and right counts
44            good_triplets_count += left_count * right_count
45            # Update the Binary Indexed Tree with the current number's position
46            tree.update(order, 1)
47
48        return good_triplets_count  # Return the total count of good triplets
49

1class Solution {
2    public long goodTriplets(int[] nums1, int[] nums2) {
3        int length = nums1.length;
4        // Array to hold the positions of elements of nums2 w.r.t nums1
5        int[] positions = new int[length];
6        // Initialize the binary indexed tree with size 'length'
7        BinaryIndexedTree bit = new BinaryIndexedTree(length);
8        // Fill the positions array with the current positions of elements from nums2
9        for (int i = 0; i < length; ++i) {
10            positions[nums2[i]] = i + 1;
11        }
12      
13        long totalGoodTriplets = 0; // This will hold the answer - the total number of good triplets
14      
15        // Traverse each element in nums1
16        for (int number : nums1) {
17            int position = positions[number];
18            // The number of elements less than the current number in nums2
19            long leftCount = bit.query(position);
20            // The number of elements greater than the current number in nums2
21            long rightCount = length - position - (bit.query(length) - bit.query(position));
22            // Multiply left count and right count to get the number of good triplets for current number
23            totalGoodTriplets += leftCount * rightCount;
24            // Update the binary indexed tree
25            bit.update(position, 1);
26        }
27        return totalGoodTriplets;
28    }
29}
30
31class BinaryIndexedTree {
32    private int size; // The size of binary indexed tree
33    private int[] tree; // The binary indexed tree data structure
34
35    public BinaryIndexedTree(int size) {
36        this.size = size;
37        tree = new int[size + 1];
38    }
39
40    // Method to update the tree with value 'delta' at index 'index'
41    public void update(int index, int delta) {
42        while (index <= size) {
43            tree[index] += delta;
44            index += lowbit(index); // Find next index to update
45        }
46    }
47
48    // Method to query the cumulative frequency up until index 'index'
49    public int query(int index) {
50        int sum = 0;
51        while (index > 0) {
52            sum += tree[index];
53            index -= lowbit(index); // Move to the previous index to query sum
54        }
55        return sum;
56    }
57
58    // Utility method to get the least significant bit
59    public static int lowbit(int value) {
60        return value & -value;
61    }
62}
63

1#include <vector>
2
3using std::vector;
4
5class BinaryIndexedTree {
6public:
7    vector<int> tree;
8    int size;
9
10    // Constructor: initializes the BIT with a specific size.
11    explicit BinaryIndexedTree(int n) : size(n), tree(n + 1, 0) {}
12
13    // Increases the value at index x by delta.
14    void update(int index, int delta) {
15        // Keep adding delta to the tree[index] and move to the next index.
16        // The next index is obtained by incrementing the last set bit of the current index.
17        while (index <= size) {
18            tree[index] += delta;
19            index += lowbit(index);
20        }
21    }
22
23    // Computes the prefix sum from index 1 to x.
24    int query(int index) const {
25        int sum = 0;
26        // Keep adding values from the tree and move to the previous index.
27        // The previous index is obtained by decrementing the last set bit of the current index.
28        while (index > 0) {
29            sum += tree[index];
30            index -= lowbit(index);
31        }
32        return sum;
33    }
34
35private:
36    // Utility function to get the last set bit of x.
37    static int lowbit(int x) {
38        return x & -x;
39    }
40};
41
42class Solution {
43public:
44    // Calculates the number of good triplets based on the given rules.
45    long long goodTriplets(const vector<int>& nums1, const vector<int>& nums2) {
46        int n = nums1.size();
47
48        // Mapping from elements to their positions in nums2.
49        vector<int> position(n);
50        for (int i = 0; i < n; ++i) position[nums2[i]] = i + 1;
51
52        // Instantiate a Binary Indexed Tree to help with the computation.
53        BinaryIndexedTree bit(n);
54        long long count = 0; // To store the final result.
55
56        // Iterate over each number in nums1.
57        for (int num : nums1) {
58            int currentPos = position[num];
59            // Query the number of elements before the current position
60            // which correspond to the left side of a good triplet.
61            int leftCount = bit.query(currentPos);
62            // Query the number of elements after the current position
63            // which correspond to the right side of a good triplet.
64            int rightCount = n - currentPos - (bit.query(n) - bit.query(currentPos));
65            // Add to the result the number of good triplets with the current number in the middle.
66            count += static_cast<long long>(leftCount) * rightCount;
67            // Mark the position of this number as seen.
68            bit.update(currentPos, 1);
69        }
70        return count;
71    }
72};
73

1// TypeScript does not make use of #include directives used in C++.
2// Also, the 'vector' term is not used in TypeScript; instead, we use arrays.
3
4// Global variable to hold the tree structure of the Binary Indexed Tree (BIT).
5let tree: number[];
6// Global variable to define the size of the BIT.
7let size: number;
8
9/**
10 * Initializes the BIT with a specific size.
11 * @param n The size of the array for which BIT is being constructed.
12 */
13function initBinaryIndexedTree(n: number): void {
14    size = n;
15    tree = new Array(n + 1).fill(0);
16}
17
18/**
19 * Increases the value at the given index by delta.
20 * @param index The index at which to increase the value.
21 * @param delta The amount to increase the value by.
22 */
23function update(index: number, delta: number): void {
24    // Keep adding delta to the tree at the specified index and update the index
25    // by incrementing it by the value of its last set bit.
26    while (index <= size) {
27        tree[index] += delta;
28        index += lowbit(index);
29    }
30}
31
32/**
33 * Computes the prefix sum from index 1 to the given index.
34 * @param index The index up to which to compute the sum.
35 * @returns The prefix sum up to the specified index.
36 */
37function query(index: number): number {
38    let sum = 0;
39    // Keep adding to the sum from the tree at the specified index and update the index
40    // by decrementing it by the value of its last set bit.
41    while (index > 0) {
42        sum += tree[index];
43        index -= lowbit(index);
44    }
45    return sum;
46}
47
48/**
49 * Utility function to get the last set bit of a number.
50 * @param x The number whose last set bit is to be found.
51 * @returns The last set bit of the given number.
52 */
53function lowbit(x: number): number {
54    return x & -x;
55}
56
57/**
58 * Calculates the number of good triplets based on the given rules.
59 * @param nums1 An array of numbers representing the first sequence.
60 * @param nums2 An array of numbers representing the second sequence.
61 * @returns The number of good triplets.
62 */
63function goodTriplets(nums1: number[], nums2: number[]): number {
64    let n = nums1.length;
65    // Map elements to their positions in nums2.
66    let position: number[] = new Array(n);
67    for (let i = 0; i < n; ++i) {
68        position[nums2[i]] = i + 1;
69    }
70
71    // Initialize the BIT for use in computing the good triplets.
72    initBinaryIndexedTree(n);
73    let count = 0; // The final result.
74
75    // Iterate over the numbers in nums1.
76    for (let num of nums1) {
77        let currentPos = position[num];
78        // The count of elements before the current position in BIT,
79        // corresponding to the left side of a good triplet.
80        let leftCount = query(currentPos);
81        // The count of elements after the current position in BIT,
82        // corresponding to the right side of a good triplet.
83        let rightCount = n - currentPos - (query(n) - query(currentPos));
84        // Increment the result by the number of good triplets with the current number in the middle.
85        count += leftCount * rightCount;
86        // Mark the position of this number as seen in the BIT.
87        update(currentPos, 1);
88    }
89    return count;
90}
91

Time and Space Complexity

Time Complexity

The given code defines a BinaryIndexedTree class and a method goodTriplets within another class Solution. Let's analyze the time complexity:

The update method of BinaryIndexedTree calls lowbit within a while loop that iterates proportional to how fast the current index x can be decremented to zero. The lowbit operation itself is O(1). The number of iterations in the worst case is O(log n), where n is the number of elements in the binary indexed tree.
The query method of BinaryIndexedTree is similar to the update method in terms of complexity, performing an iterative process with a lowbit operation, and also has a time complexity of O(log n).
The goodTriplets method initializes a BinaryIndexedTree, which takes O(n) time because a list of n + 1 zeros is created.
Within the goodTriplets method, a loop iterates over each element of nums1. Inside this loop, query and update operations of BinaryIndexedTree are performed, each taking O(log n) time.
For each element in nums1, two query calls and one update call are made, resulting in O(log n) time complexity per element.
Consequently, the loop in goodTriplets runs in O(n log n) time, as there are n elements to process and each element requires O(log n) time.

Therefore, the overall time complexity of the code is O(n log n).

Space Complexity

The BinaryIndexedTree class consumes space for an array of length n + 1, giving us a space complexity of O(n).
The dictionary pos is created with n key-value pairs, also contributing to the space complexity of O(n).
Local variables within the method goodTriplets contribute an insignificant constant space overhead.

Combining these factors, the overall space complexity of the code is O(n).

Learn more about how to find time and space complexity quickly using problem constraints.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Question 1 out of 10

Which algorithm should you use to find a node that is close to the root of the tree?

Depth First Search

Breadth First Search

I don’t know