Leetcode 2179. Count Good Triplets in an Array

Jimmy   Apr 15, 2025   1316 words   ~35 mins
#Array #Binary Search #Divide and Conquer #Binary Indexed Tree #Segment Tree #Merge Sort #Ordered Set

Table of Contents

Problem Informations

  • Problem Index: 2179
  • Problem Link: Count Good Triplets in an Array
  • Topics: Array, Binary Search, Divide and Conquer, Binary Indexed Tree, Segment Tree, Merge Sort, Ordered Set

Problem Description

You are given two 0-indexed arrays nums1 and nums2 of length $n$, both of which are permutations of $[0, 1, \ldots, 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 \leq x, y, z \leq 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 $pos1_x < pos1_y < pos1_z$. 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 $pos2_x < pos2_y < pos2_z$. 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 \leq n \leq 10^5$
  • $0 \leq nums1[i], nums2[i] \leq n - 1$
  • nums1 and nums2 are permutations of $[0, 1, \ldots, n - 1]$.

Intuition

The problem requires us to determine the number of “good triplets” in two permutations of a sequence. A good triplet consists of three elements that appear in a specific increasing order in both permutations. Given the constraint that nums1 and nums2 are permutations of the same set, we can leverage this property to track relative positions efficiently, enabling us to identify valid triplets without needing to directly compare them in an $O(n^3)$ fashion.

Approach

The given solution utilizes a data structure known as a Fenwick Tree (also called a Binary Indexed Tree), which is effective for querying prefix sums over an array and performing point updates in logarithmic time. This is particularly suitable for addressing the counting of inversions or, in our case, counting valid triplet configurations.

  1. Mapping Positions:
    First, a mapping of the positions of elements in nums2 is created. This is stored in an array pos2 where pos2[val] gives the position of val in the nums2 array. This mapping helps in transition from the order in nums1 to the order in nums2.

  2. Index Mapping from nums1 to nums2:
    Using the pos2 mapping, the elements in nums1 are transformed into their respective indices in nums2, resulting in indexMapping. Here, indexMapping[i] gives the position of the element nums1[i] in nums2.

  3. Fenwick Tree Initialization:
    A Fenwick Tree is initialized to keep track of the prefix sums efficiently. This structure supports logarithmic updates and prefix queries.

  4. Counting Good Triplets:
    Iterate through each position in indexMapping, treating this as the middle element mid of a potential triplet:

    • Calculate left: the number of indices smaller than the current mid that hold a smaller value, which can be found using a prefix sum query up to indexMapping[mid] in the Fenwick Tree.
    • Calculate right: the potential larger values to the right of the current mid, computed as the remaining indices to the right minus those counted as left.
    • Update the Fenwick Tree to include the current position.
    • The number of good triplets that can end at the current mid is given by multiplying left and right because these are independent choices.

  5. Complexity Consideration:
    Each update and query operation in the Fenwick Tree is $O(\log n)$, and the overall approach involves a single pass through the array, making the complexity approximately $O(n \log n)$. This is efficient given the problem’s constraints.

By following these steps, the algorithm effectively counts the number of good triplets, adhering to both the conditions specified in the problem statement.

Code

C++

class FenwickTree {
private:
    vector<int> tree;

    // Return the lowest bit (1) of the number
    int lowbit(int num) { 
        return num & -num; 
    }

public:
    // Initialize the Fenwick Tree with given size
    FenwickTree(int size) { 
        tree = vector(size + 1, 0); 
    }

    // Update the Fenwick Tree at a specific index with a delta value
    void update(int idx, int delta) {
        idx++;
        while (idx < tree.size()) {
            tree[idx] += delta;
            idx += lowbit(idx);
        }
    }

    // Query the prefix sum up to a given index in the Fenwick Tree
    int query(int idx) {
        idx++;
        int sum = 0;
        while (idx > 0) {
            sum += tree[idx];
            idx -= lowbit(idx);
        }
        return sum;
    }
};

class Solution {
public:
    long long goodTriplets(vector<int>& nums1, vector<int>& nums2) {
        int n = nums1.size();
        long long totalGoodTriplets = 0;
        vector<int> pos2(n), indexMapping(n);

        // Record positions of each element in nums2
        for (int i = 0; i < n; i++) {
            pos2[nums2[i]] = i;
        }

        // Map nums1 elements' positions in nums2
        for (int i = 0; i < n; i++) {
            indexMapping[i] = pos2[nums1[i]];
        }

        FenwickTree tree(n);

        // Calculate good triplets
        for (int mid = 0; mid < n; mid++) {
            int pos = indexMapping[mid];
            int left = tree.query(pos);
            int right = (n - pos - 1) - (mid - left);

            // Update Fenwick Tree for the current mid position
            tree.update(pos, 1);

            // Calculate number of good triplets based on the current left and right values
            totalGoodTriplets += static_cast<long long>(left) * right;
        }

        return totalGoodTriplets;
    }
};

Python

class FenwickTree:
    def __init__(self, size):
        # Initialize the Fenwick Tree with given size
        self.tree = [0] * (size + 1)

    def lowbit(self, num):
        # Return the lowest bit (1) of the number
        return num & -num

    def update(self, idx, delta):
        # Update the Fenwick Tree at a specific index with a delta value
        idx += 1
        while idx < len(self.tree):
            self.tree[idx] += delta
            idx += self.lowbit(idx)

    def query(self, idx):
        # Query the prefix sum up to a given index in the Fenwick Tree
        idx += 1
        sum = 0
        while idx > 0:
            sum += self.tree[idx]
            idx -= self.lowbit(idx)
        return sum

class Solution:
    def goodTriplets(self, nums1, nums2):
        n = len(nums1)
        totalGoodTriplets = 0
        pos2 = [0] * n
        indexMapping = [0] * n

        # Record positions of each element in nums2
        for i in range(n):
            pos2[nums2[i]] = i

        # Map nums1 elements' positions in nums2
        for i in range(n):
            indexMapping[i] = pos2[nums1[i]]

        tree = FenwickTree(n)

        # Calculate good triplets
        for mid in range(n):
            pos = indexMapping[mid]
            left = tree.query(pos)
            right = (n - pos - 1) - (mid - left)

            # Update Fenwick Tree for the current mid position
            tree.update(pos, 1)

            # Calculate number of good triplets based on the current left and right values
            totalGoodTriplets += left * right

        return totalGoodTriplets

Complexity

  • Time complexity: $O(n \log n)$

    The algorithm primarily consists of three main components which determine its time complexity:

    1. Initial Setup: It takes $O(n)$ time to fill pos2 and indexMapping arrays, where $n$ is the length of nums1 (or nums2).

    2. Fenwick Tree Updates and Queries: The main loop iterates over each element in indexMapping, and within each iteration, it performs a query and an update operation on the Fenwick Tree. Both the update and query operations on a Fenwick Tree take $O(\log n)$ time. Since this loop runs $n$ times, the overall time complexity for this part is $O(n \log n)$.

    Combining these components, the dominant term is $O(n \log n)$ due to the Fenwick Tree operations, which makes the overall time complexity $O(n \log n)$.

  • Space complexity: $O(n)$

    The space complexity is primarily determined by the additional storage used:

    1. Fenwick Tree Storage: The Fenwick Tree itself maintains an array tree of size $n+1$.

    2. Auxiliary Arrays: Two auxiliary arrays pos2 and indexMapping are each of size $n$.

    Therefore, the overall space complexity is $O(n)$, as the size of the Fenwick Tree and the auxiliary arrays dominate the space usage.

Disclaimer: All reference materials on this website are sourced from the internet and are intended for learning purposes only. If you believe any content infringes upon your rights, please contact me at csnote.cc@gmail.com, and I will remove the relevant content promptly.


Feedback Welcome: If you notice any errors or areas for improvement in the articles, I warmly welcome your feedback and corrections. Your input will help this blog provide better learning resources. This is an ongoing process of learning and improvement, and your suggestions are valuable to me. You can reach me at csnote.cc@gmail.com.