2179. Count Good Triplets in an Array

Description

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

Solution

count-good-triplets-in-an-array.py

from sortedcontainers import SortedList

class Solution:
    def goodTriplets(self, nums1: List[int], nums2: List[int]) -> int:
        n = len(nums1)
        pos = {x:i for i,x in enumerate(nums2)}

        sl = SortedList([pos[nums1[0]]])
        prefix = [0]

        for i in range(1, n):
            sl.add(pos[nums1[i]])
            prefix.append(sl.bisect_left(pos[nums1[i]]))

        sl = SortedList([pos[nums1[-1]]])
        suffix = [0]

        for i in range(n - 2, -1, -1):
            index = sl.bisect_left(pos[nums1[i]])
            suffix.append(len(sl) - index)
            sl.add(pos[nums1[i]])
        suffix.reverse()

        res = 0
        for a, b in zip(prefix, suffix):
            res += a * b

        return res