Count Good Triplets in an Array LeetCode solution in C++

1
typedef long long ll;
2
typedef long double ld;
3
#define mod 1000000007
4
#define all(x) begin(x), end(x)
5
#include <ext/pb_ds/assoc_container.hpp>
6
#include <ext/pb_ds/tree_policy.hpp>
7
using namespace __gnu_pbds;
8
#define ordered_set tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update>
9

10
class Solution {
11
 public:
12
  long long goodTriplets(vector<int> &nums1, vector<int> &nums2) {
13
    unordered_map<int, int> mapping;
14
    int no = 1;
15
    for (int x : nums1) {
16
      mapping[x] = no;
17
      no++;
18
    }
19
    for (int i = 0; i < nums2.size(); ++i) nums2[i] = mapping[nums2[i]];
20
    int n = nums2.size();
21
    ordered_set st1;
22
    ordered_set st2;
23
    for (int x : nums2) st2.insert(x);
24
    ll ans = 0;
25
    st1.insert(nums2[0]);
26
    st2.erase(st2.find(nums2[0]));
27
    for (int i = 1; i < n - 1; ++i) {
28
      st2.erase(st2.find(nums2[i]));
29
      ll less = (ll)st1.order_of_key(nums2[i]);
30
      ll great = (ll)(((int)st2.size()) - st2.order_of_key(nums2[i]));
31
      ans += (less * great);
32
      st1.insert(nums2[i]);
33
    }
34
    return ans;
35
  }
36
};