LeetCode in Kotlin

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

Solution

@Suppress("NAME_SHADOWING")
class Solution {
    fun goodTriplets(nums1: IntArray, nums2: IntArray): Long {
        val n = nums1.size
        val idx = IntArray(n)
        val arr = IntArray(n)
        for (i in 0 until n) {
            idx[nums2[i]] = i
        }
        for (i in 0 until n) {
            arr[i] = idx[nums1[i]]
        }
        val tree = Tree(n)
        var res = 0L
        for (i in 0 until n) {
            val smaller = tree.query(arr[i])
            val bigger = n - (arr[i] + 1) - (i - smaller)
            res += smaller.toLong() * bigger
            tree.update(arr[i] + 1, 1)
        }
        return res
    }

    private class Tree(var n: Int) {
        var array: IntArray

        init {
            array = IntArray(n + 1)
        }

        fun lowbit(x: Int): Int {
            return x and -x
        }

        fun update(i: Int, delta: Int) {
            var i = i
            while (i <= n) {
                array[i] += delta
                i += lowbit(i)
            }
        }

        fun query(k: Int): Int {
            var k = k
            var ans = 0
            while (k > 0) {
                ans += array[k]
                k -= lowbit(k)
            }
            return ans
        }
    }
}

This site is open source. Improve this page.