LeetCode Solutions Blog

Count Good Triplets in an Array

Number: 2280

Difficulty: Hard

Paid? No

Problem Description

Given two 0-indexed arrays nums1 and nums2 which are permutations of [0, 1, ..., n - 1], we need to count the total number of good triplets. A good triplet is defined as a set of three distinct values (x, y, z) such that their positions are in increasing order in both arrays. In terms of indices, if pos1_v is the index of value v in nums1 and pos2_v is the index of value v in nums2, then a triplet (x, y, z) is good if and only if pos1_x < pos1_y < pos1_z and pos2_x < pos2_y < pos2_z.

Key Insights

Since nums1 and nums2 are permutations, a mapping from values to their positions can be built.
Transform the problem by mapping nums1 into an array P where P[i] = pos2[nums1[i]]. Then, counting good triplets is equivalent to counting increasing triplets (i, j, k) in array P.
For each element at position j in P, count:
- left_count: the number of indices i < j with P[i] < P[j]
- right_count: the number of indices k > j with P[k] > P[j]
Use a Binary Indexed Tree (BIT) or Fenwick Tree to efficiently count elements in a given range.
Final answer is the sum over all j of left_count[j] * right_count[j].

Space and Time Complexity

Time Complexity: O(n log n)
Space Complexity: O(n)

Solution

We first convert the problem into counting increasing triplets in a transformed array P, where each element is the position of nums1[i] in nums2. We then use two BITs:

One BIT (leftBIT) processes from left to right, counting for each position j the number of previous elements that are less than P[j].
Another BIT (rightBIT) processes from right to left, counting for each position j the number of later elements that are greater than P[j]. For each index j in the array P, the product left_count[j] * right_count[j] gives the number of good triplets with the middle element at j. The sum of these products over all j is the answer.

The BIT provides O(log n) update and query operations, ensuring the overall solution remains efficient even for n up to 10^5.

Code Solutions

# Python solution with line-by-line comments

class FenwickTree:
    # Binary Indexed Tree for sum queries
    def __init__(self, size):
        self.size = size
        self.tree = [0] * (size + 1)
    
    def update(self, index, value):
        # Update the BIT at position (index+1) with value
        index += 1  # BIT uses 1-indexed positions
        while index <= self.size:
            self.tree[index] += value
            index += index & -index
    
    def query(self, index):
        # Query cumulative sum from start to index (inclusive)
        result = 0
        index += 1  # adjust to BIT indexing
        while index:
            result += self.tree[index]
            index -= index & -index
        return result

def count_good_triplets(nums1, nums2):
    n = len(nums1)
    # map each number v to its index in nums2
    pos_in_nums2 = [0] * n
    for index, value in enumerate(nums2):
        pos_in_nums2[value] = index
    
    # Create transformed array P where P[i] = pos of nums1[i] in nums2
    P = [pos_in_nums2[value] for value in nums1]
    
    # Initialize arrays to hold counts for left and right side
    left_counts = [0] * n
    right_counts = [0] * n
    
    # BIT for counting elements on the left side
    leftBIT = FenwickTree(n)
    for i in range(n):
        # Number of elements with value < P[i]
        left_counts[i] = leftBIT.query(P[i] - 1)
        leftBIT.update(P[i], 1)
    
    # BIT for counting elements on the right side
    rightBIT = FenwickTree(n)
    for i in range(n - 1, -1, -1):
        # Number of elements with value > P[i]
        # Total count on right side so far minus count of numbers <= P[i]
        right_counts[i] = rightBIT.query(n - 1) - rightBIT.query(P[i])
        rightBIT.update(P[i], 1)
    
    # Count valid triplets by multiplying counts from left and right for each middle element
    result = 0
    for i in range(n):
        result += left_counts[i] * right_counts[i]
    return result

# Example usage:
nums1 = [2, 0, 1, 3]
nums2 = [0, 1, 2, 3]
print(count_good_triplets(nums1, nums2))  # Output: 1