## Problem

Problem ID: 1511

Title: Count Number of Teams

Difficulty: Medium

Description:

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).

Example 1:

Input: rating = [2,5,3,4,1]
Output: 3
Explanation: We can form three teams given the conditions. (2,3,4), (5,4,1), (5,3,1).


Example 2:

Input: rating = [2,1,3]
Output: 0
Explanation: We can't form any team given the conditions.


Example 3:

Input: rating = [1,2,3,4]
Output: 4


Constraints:

• n == rating.length
• 3 <= n <= 1000
• 1 <= rating[i] <= 105
• All the integers in rating are unique.

## Thoughts

As usual, I didn’t take a step back to think about the most optimal solution before attempting to solve this problem. The initial approach I took was the usual DP counting problem where the number of ways to form a team of 2 is add the number of ways soldier i can join a team of 1 and the number of ways to form a team of 3 is to add the number of ways soldier i can join a team of 2.

Initial solution (O(N^2)):

class Solution {
public:
int helper(const vector<int>& rating) {
const auto n = rating.size();
vector<int> dp(n,1);
for (size_t soldier = 1; soldier < 3; soldier++) {
vector<int> next_dp(n, 0);
for (size_t j = 0; j < n; j++) {
for (size_t i = 0; i < j; i++) {
if (rating[i] >= rating[j]) continue;
next_dp[j] += dp[i];
}
}
dp = move(next_dp);
}
int ret = 0;
for (const auto count : dp) {
ret += count;
}
return ret;
}
int numTeams(vector<int>& rating) {
int forward = helper(rating);
reverse(rating.begin(), rating.end());
int backward = helper(rating);
return forward + backward;
}
};


## Solution

General Idea

Intuition: Treat each solider as the middle j solider and find the number of possible of i (left) and number of possible k (right).

To find the number of possible i and k:

• Two multiset (sorted)
• One multiset represents the soldiers to the left of j
• One represents the soldiers to the right of j.
• Finding the number of possible i is the number of elements in the left multiset that are strictly less than rating[j]
• This can be easily achieved by using lower_bound (the first element that >= rating[j]) and moving it back by 1
• Use std::distance to find the distance between the start and the new iterator
• Finding the number of possible j is the number of elements in the right multiset that are strictly greater than rating[j]
• This can be easily achieved by using upper_bound (the first element that > rating[j])
• Use std::distance to find the distance between the new iterator and the end
• Complexity: O(log(N))

Multiplying the number of possible i and k will give the total groups with solider j in the middle.

### Implementation

class Solution {
public:
int solve(const vector<int>& ratings) {
const auto n = ratings.size();
multiset<int> left;
multiset<int> right;
for (const auto rating : ratings) {
right.insert(rating);
}
int ret = 0;
for (size_t i = 0; i < n; i++) {
auto rating = ratings[i];
right.erase(right.find(rating));

auto l_it = lower_bound(left.begin(), left.end(), rating);
if (l_it !=  left.begin()) {
// the largest left that is strictly less than rating
l_it--;
auto left_count = distance(left.begin(), l_it) +1;
// the smallest right that is strictly more than rating
auto r_it = upper_bound(right.begin(), right.end(), rating);
auto right_count = distance(r_it, right.end());
ret += (left_count*right_count);
}
left.insert(rating);
}
return ret;
}
int numTeams(vector<int>& rating) {
int forward = solve(rating);
reverse(rating.begin(), rating.end());
int backward = solve(rating);
return forward + backward;

}
};