1395. Count Number of Teams

There are n soldiers standing in a line. Each soldier is assigned a unique rating value.

You have to form a team of 3 soldiers amongst them under the following rules:

• Choose 3 soldiers with index (i, j, k) with rating (rating[i], rating[j], rating[k]).
• A team is valid if: (rating[i] < rating[j] < rating[k]) or (rating[i] > rating[j] > rating[k]) where (0 <= i < j < k < n).

Return the number of teams you can form given the conditions. (soldiers can be part of multiple teams).

c++

class Solution {
public:
    int numTeams(vector<int>& r) {
        set<int> left(begin(r), next(begin(r))), right(next(begin(r)), end(r));
        int res(0), sz = r.size();
        for(int i = 1; i < sz - 1; i++) {
            right.erase(r[i]);
            int leftGreater = distance(left.lower_bound(r[i]), end(left)), rightGreater = distance(right.lower_bound(r[i]), end(right)), leftLess = left.size() - leftGreater, rightLess = right.size() - rightGreater;
            res += leftGreater * rightLess + leftLess * rightGreater;
            left.insert(r[i]);
        }
        return res;
    }
};