2179. Count Good Triplets in an Array

2179. Count Good Triplets in an Array (Hard)

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 pos1v as the index of the value v in nums1 and pos2v 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 pos1x < pos1y < pos1z and pos2x < pos2y < pos2z.

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 <= 105
• 0 <= nums1[i], nums2[i] <= n - 1
• nums1 and nums2 are permutations of [0, 1, ..., n - 1].

Similar Questions:

• Count of Smaller Numbers After Self (Hard)
• Increasing Triplet Subsequence (Medium)
• Create Sorted Array through Instructions (Hard)

Solution 1. Divide and Conquer (Merge Sort)

The first idea is that we pick a number as the middle number in the triplet, and count the common numbers in front of this number and after this number.

For example,

A = [4,0,1,3,2]
B = [4,1,0,2,3]

For number 1, there is a single common number (4) in front of 1 and two common numbers (3,4) after 1, so the count of triplets with 1 in the middle is 1 * 2 = 2.

But counting the common numbers is not easy. Since the numbers are permutations of [0, N-1], we can simplify the problem by mapping on of the array to [0, N-1].

For example, if we map A to [0, N-1].

A = [4,0,1,3,2]
A'= [0,1,2,3,4]
mapping = 4->0, 0->1, 1->2, 3->3, 2->4

We apply the same mapping to B

B = [4,1,0,2,3]
B'= [0,2,1,4,3]

Now look at the new arrays

A = [0,1,2,3,4]
B = [0,2,1,4,3]

For the number 1, we trivially know that in A, there is one number 0 before it and 3 numbers 2,3,4 after it. So, we just need to look at B. The problem now becomes counting the numbers smaller than 1 before 1 and numbers greater than 1 after 1.

Let lt[i] be the count of numbers smaller than B[i]. The count of common numbers before B[i] is min(B[i], lt[i]) = lt[i]. The count of common numbers after B[i] in A is N - B[i] - 1. The count of common numbers after B[i] in B is N - i - 1 - (B[i] - lt[i]) = N - B[i] - 1 - i + lt[i]. Since i >= lt[i], -i + lt[i] <= 0, the count of common numbers after B[i] in both arrays is N - B[i] - 1 - i + lt[i].

So, the count of triplets with B[i] as the middle number is min(B[i], lt[i]) * (N - B[i] - 1 - i + lt[i])

// OJ: https://leetcode.com/problems/count-good-triplets-in-an-array/
// Author: github.com/lzl124631x
// Time: O(NlogN)
// Space: O(N)
class Solution {
public:
    long long goodTriplets(vector<int>& A, vector<int>& B) {
        long N = A.size(), ans = 0;
        vector<int> m(N), lt(N), tmp(N), tmpLt(N), index(N);
        for (int i = 0; i < N; ++i) m[A[i]] = i;
        for (int i = 0; i < N; ++i) {
            B[i] = m[B[i]];
            index[B[i]] = i;
        }
        function<void(int, int)> merge = [&](int begin, int end) {
            if (begin + 1 >= end) return;
            int mid = (begin + end) / 2;
            merge(begin, mid);
            merge(mid, end);
            int i = begin, j = mid, k = begin;
            for (; k < end; ++k) {
                if (j >= end || (i < mid && B[i] < B[j])) {
                    tmp[k] = B[i];
                    tmpLt[k] = lt[i];
                    ++i;
                } else {
                    tmp[k] = B[j];
                    tmpLt[k] = lt[j] + i - begin;
                    ++j;
                }
            }
            for (int i = begin; i < end; ++i) {
                B[i] = tmp[i];
                lt[i] = tmpLt[i];
            }
        };
        merge(0, N);
        for (int i = 0; i < N; ++i) {
            ans += (long)lt[i] * (N - B[i] - 1 - index[B[i]] + lt[i]);
        }
        return ans;
    }
};

Solution 2. Binary Index Tree

// OJ: https://leetcode.com/problems/count-good-triplets-in-an-array/
// Author: github.com/lzl124631x
// Time: O(NlogN)
// Space: O(N)
// Ref: https://leetcode-cn.com/problems/count-good-triplets-in-an-array/solution/deng-jie-zhuan-huan-shu-zhuang-shu-zu-by-xmyd/
class Solution {
public:
    long long goodTriplets(vector<int> &A, vector<int> &B) {
        int N = A.size();
        vector<int> m(N);
        for (int i = 0; i < N; ++i) m[A[i]] = i;
        long long ans = 0;
        vector<int> tree(N + 1);
        for (int i = 1; i < N - 1; ++i) {
            for (int j = m[B[i - 1]] + 1; j <= N; j += j & -j) ++tree[j]; // increment the count of m[B[i-1]]
            int y = m[B[i]], less = 0;
            for (int j = y; j; j &= j - 1) less += tree[j]; // count all the numbers less than m[B[i]]
            ans += (long)less * (N - 1 - y - (i - less));
        }
        return ans;
    }
};