Count Good Triplets in an Array LeetCode solution in JavaScript

1
var goodTriplets = function (nums1, nums2) {
2
  let n = nums1.length,
3
    nums2_idx = Array(n);
4
  for (let i = 0; i < n; i++) nums2_idx[nums2[i]] = i;
5
  let idxs = Array(n);
6
  for (let i = 0; i < n; i++) idxs[i] = nums2_idx[nums1[i]];
7

8
  let smallerLeft = getSmallerLeft(idxs); // smallerLeft[i] = number of indices to the left of i smaller than idxs[i] in both nums1 and nums2
9
  let biggerRight = getBiggerRight(idxs); // biggerRight[i] = number of indices to the right of i bigger than idxs[i] in both nums1 and nums2
10
  let ans = 0;
11
  for (let i = 1; i < n - 1; i++) {
12
    ans += smallerLeft[i] * biggerRight[i];
13
  }
14
  return ans;
15
};
16

17
function getSmallerLeft(idxs) {
18
  let n = idxs.length,
19
    res = Array(n).fill(0);
20
  let segTree = new SegmentTree(n);
21
  for (let i = 0; i < n; i++) {
22
    res[i] = segTree.getSum(0, idxs[i] - 1);
23
    segTree.update(idxs[i], 1);
24
  }
25
  return res;
26
}
27

28
function getBiggerRight(idxs) {
29
  let n = idxs.length,
30
    res = Array(n).fill(0);
31
  let segTree = new SegmentTree(n);
32
  for (let i = n - 1; i >= 0; i--) {
33
    res[i] = segTree.getSum(idxs[i] + 1, n - 1);
34
    segTree.update(idxs[i], 1);
35
  }
36
  return res;
37
}
38

39
class SegmentTree {
40
  constructor(n) {
41
    this.size = n;
42
    this.segTree = Array(n * 2).fill(0);
43
  }
44
  update(index, val) {
45
    let n = this.size,
46
      idx = index + n;
47
    this.segTree[idx] += val;
48
    idx = Math.floor(idx / 2);
49

50
    while (idx > 0) {
51
      this.segTree[idx] = this.segTree[idx * 2] + this.segTree[idx * 2 + 1];
52
      idx = Math.floor(idx / 2);
53
    }
54
  }
55
  getSum(left, right) {
56
    let n = this.size,
57
      sum = 0;
58
    let left_idx = left + n,
59
      right_idx = right + n;
60
    // left must be even, right must be odd
61
    while (left_idx <= right_idx) {
62
      if (left_idx % 2 === 1) sum += this.segTree[left_idx++];
63
      if (right_idx % 2 === 0) sum += this.segTree[right_idx--];
64
      left_idx = Math.floor(left_idx / 2);
65
      right_idx = Math.floor(right_idx / 2);
66
    }
67
    return sum;
68
  }
69
}