Count Good Triplets in an Array LeetCode solution in Python

1
from sortedcontainers import SortedList
2

3

4
class Solution:
5
    def goodTriplets(self, A: List[int], B: List[int]) -> int:
6
        # Index of a (from A) in B.
7
        pos = [0] * len(A)
8
        for idx, b in enumerate(B):
9
            pos[b] = idx
10

11
        # Build pre_a[i]: number of elements on a[i]'s left in both A and B.
12
        # pos_in_b: sorted indexes (in B) of all the visited elements in A.
13
        pos_in_b, pre_a = SortedList([pos[A[0]]]), [0]
14
        for a in A[1:]:
15
            pos_in_b.add(pos[a])
16
            pre_a.append(pos_in_b.bisect_left(pos[a]))
17

18
        # Build suf_a[i]: number of elements on a[i]'s right in both A and B.
19
        pos_in_b, suf_a = SortedList([pos[A[-1]]]), [0]
20
        for a in reversed(A[: len(A) - 1]):
21
            idx = pos_in_b.bisect(pos[a])
22
            suf_a.append(len(pos_in_b) - idx)
23
            pos_in_b.add(pos[a])
24
        suf_a.reverse()
25

26
        # Sum up all unique triplets centered on A[i].
27
        ans = 0
28
        for x, y in zip(pre_a, suf_a):
29
            ans += x * y
30
        return ans