Count Number of Teams
MedPrompt
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.length3 <= n <= 10001 <= rating[i] <= 105- All the integers in
ratingare unique.
Approaches
3 approaches with complexity analysis and trade-offs.
This is the most straightforward approach. We can generate all possible triplets of soldiers (i, j, k) such that 0 <= i < j < k < n using three nested loops. For each triplet, we check if their ratings satisfy the condition for a valid team: rating[i] < rating[j] < rating[k] (increasing) or rating[i] > rating[j] > rating[k] (decreasing). If the condition is met, we increment a counter.
Algorithm
- Initialize a counter
teamsto 0. - Use a loop for the first soldier
ifrom0ton-3. - Inside this loop, use a nested loop for the second soldier
jfromi+1ton-2. - Inside the second loop, use another nested loop for the third soldier
kfromj+1ton-1. - Within the innermost loop, check if
rating[i] < rating[j] && rating[j] < rating[k]. If true, incrementteams. - Also, check if
rating[i] > rating[j] && rating[j] > rating[k]. If true, incrementteams. - After all loops complete, return
teams.
Walkthrough
The brute-force method systematically checks every possible combination of three distinct soldiers. It uses three nested loops to select indices i, j, and k ensuring that i < j < k. For each valid combination of indices, it fetches their ratings and checks if they form a strictly increasing or strictly decreasing sequence. A counter is maintained to keep track of the number of valid teams found.
class Solution { public int numTeams(int[] rating) { int n = rating.length; if (n < 3) { return 0; } int teams = 0; for (int i = 0; i < n; i++) { for (int j = i + 1; j < n; j++) { for (int k = j + 1; k < n; k++) { if (rating[i] < rating[j] && rating[j] < rating[k]) { teams++; } if (rating[i] > rating[j] && rating[j] > rating[k]) { teams++; } } } } return teams; }}Complexity
Time
O(n^3). Three nested loops iterate through all possible triplets, where `n` is the number of soldiers. For `n=1000`, this would be approximately `10^9` operations, which is too slow for typical time limits.
Space
O(1). We only use a constant amount of extra space for the counter and loop variables.
Trade-offs
Pros
Simple to understand and implement.
Cons
Very inefficient and will likely result in a "Time Limit Exceeded" error for the given constraints (n <= 1000).
We can improve the time complexity by changing our perspective. Instead of iterating through all triplets, we can iterate through each soldier j and consider them as the middle element of a team. For each j, we need to find how many soldiers i to its left (i < j) have a smaller rating and how many have a larger rating. Similarly, we need to find how many soldiers k to its right (k > j) have a larger rating and how many have a smaller rating.
Algorithm
- Initialize a counter
teamsto 0. - Iterate through each soldier
jfrom1ton-2to be the middle element. - For each
j, initializeleftLess = 0,leftGreater = 0,rightLess = 0,rightGreater = 0. - Iterate
ifrom0toj-1to count elements on the left:- If
rating[i] < rating[j], incrementleftLess. - If
rating[i] > rating[j], incrementleftGreater.
- If
- Iterate
kfromj+1ton-1to count elements on the right:- If
rating[k] < rating[j], incrementrightLess. - If
rating[k] > rating[j], incrementrightGreater.
- If
- Add
(leftLess * rightGreater) + (leftGreater * rightLess)to the totalteams. - After the main loop finishes, return
teams.
Walkthrough
For each soldier j at index 1 to n-2:
- Count soldiers
ito the left (i < j) withrating[i] < rating[j]. Let this beleftLess. - Count soldiers
ito the left (i < j) withrating[i] > rating[j]. Let this beleftGreater. - Count soldiers
kto the right (k > j) withrating[k] > rating[j]. Let this berightGreater. - Count soldiers
kto the right (k > j) withrating[k] < rating[j]. Let this berightLess.
The number of increasing teams with j as the middle element is leftLess * rightGreater.
The number of decreasing teams with j as the middle element is leftGreater * rightLess.
We sum these products for all possible middle elements j to get the total number of teams.
class Solution { public int numTeams(int[] rating) { int n = rating.length; if (n < 3) { return 0; } int teams = 0; for (int j = 1; j < n - 1; j++) { int leftLess = 0, leftGreater = 0; for (int i = 0; i < j; i++) { if (rating[i] < rating[j]) { leftLess++; } else if (rating[i] > rating[j]) { leftGreater++; } } int rightLess = 0, rightGreater = 0; for (int k = j + 1; k < n; k++) { if (rating[k] < rating[j]) { rightLess++; } else if (rating[k] > rating[j]) { rightGreater++; } } teams += (leftLess * rightGreater) + (leftGreater * rightLess); } return teams; }}Complexity
Time
O(n^2). The outer loop runs `n-2` times. Inside it, we have two separate loops that, in total, iterate through the rest of the `n-1` elements. This results in a quadratic time complexity, which is efficient enough for `n <= 1000`.
Space
O(1). We only use a few variables to store counts, requiring constant extra space.
Trade-offs
Pros
Significantly more efficient than the brute-force approach.
Passes the time limits for the given constraints.
Simple logic and easy to implement without complex data structures.
Cons
While efficient enough for the given constraints, it is not the most optimal solution possible.
The O(n^2) approach can be optimized further. The bottleneck is repeatedly counting smaller/larger elements on the left and right for each middle element j. This counting process can be accelerated using a data structure like a Fenwick Tree (also known as a Binary Indexed Tree or BIT). A BIT allows us to find the number of elements in a certain range and update element counts in logarithmic time.
Algorithm
- Define the maximum possible rating
maxRating(e.g., 100001 based on constraints). - Create arrays
leftLess,leftGreater,rightLess,rightGreaterof sizen. - Calculate left counts:
- Initialize a Fenwick Tree (BIT)
leftBitof sizemaxRating. - Iterate
ifrom0ton-1:r = rating[i]leftLess[i] = leftBit.query(r - 1)leftGreater[i] = leftBit.query(maxRating - 1) - leftBit.query(r)leftBit.update(r, 1)
- Initialize a Fenwick Tree (BIT)
- Calculate right counts:
- Initialize a BIT
rightBitof sizemaxRating. - Iterate
ifromn-1down to0:r = rating[i]rightLess[i] = rightBit.query(r - 1)rightGreater[i] = rightBit.query(maxRating - 1) - rightBit.query(r)rightBit.update(r, 1)
- Initialize a BIT
- Calculate total teams:
- Initialize
teams = 0. - Iterate
ifrom0ton-1:teams += leftLess[i] * rightGreater[i] + leftGreater[i] * rightLess[i]
- Initialize
- Return
teams.
Walkthrough
The core idea is to pre-calculate the counts of smaller and larger elements to the left and right for every soldier. This can be done efficiently in two passes using a Fenwick Tree.
-
Pass 1 (Left to Right): We iterate from left to right. We use a BIT to store the frequencies of ratings encountered so far. For each soldier
i, we can query the BIT to find how many soldiers to its left have a smaller rating (leftLess[i]) and how many have a larger rating (leftGreater[i]). After querying, we update the BIT withrating[i]. -
Pass 2 (Right to Left): We do a similar pass from right to left with a second BIT to compute
rightLess[i]andrightGreater[i]for every soldieri. -
Final Calculation: With all four counts pre-computed for each soldier, we can iterate one last time and sum up the products
(leftLess[i] * rightGreater[i])for increasing teams and(leftGreater[i] * rightLess[i])for decreasing teams.
class Solution { public int numTeams(int[] rating) { int n = rating.length; int maxRating = 100001; int[] leftLess = new int[n]; int[] leftGreater = new int[n]; FenwickTree leftBit = new FenwickTree(maxRating); for (int i = 0; i < n; i++) { int r = rating[i]; leftLess[i] = leftBit.query(r - 1); leftGreater[i] = leftBit.query(maxRating - 1) - leftBit.query(r); leftBit.update(r, 1); } int[] rightLess = new int[n]; int[] rightGreater = new int[n]; FenwickTree rightBit = new FenwickTree(maxRating); for (int i = n - 1; i >= 0; i--) { int r = rating[i]; rightLess[i] = rightBit.query(r - 1); rightGreater[i] = rightBit.query(maxRating - 1) - rightBit.query(r); rightBit.update(r, 1); } int teams = 0; for (int i = 0; i < n; i++) { teams += leftLess[i] * rightGreater[i]; // Increasing teams += leftGreater[i] * rightLess[i]; // Decreasing } return teams; }} class FenwickTree { private int[] bit; private int size; public FenwickTree(int size) { this.size = size; this.bit = new int[size]; } public void update(int index, int delta) { while (index < size) { bit[index] += delta; index += index & -index; } } public int query(int index) { int sum = 0; while (index > 0) { sum += bit[index]; index -= index & -index; } return sum; }}Complexity
Time
O(n log M), where `n` is the number of soldiers and `M` is the maximum possible rating. We perform two passes over the `rating` array. In each pass, for each element, we perform a query and an update on the BIT, both of which take O(log M) time.
Space
O(n + M), where `n` is the number of soldiers and `M` is the maximum possible rating. We use arrays of size `n` to store the counts (O(n)) and two BITs of size `M` (O(M)).
Trade-offs
Pros
The most efficient solution with a time complexity better than quadratic.
Cons
More complex to implement due to the requirement of a Fenwick Tree or similar data structure.
Uses more space than the O(n^2) approach due to the BIT and count arrays.
Solutions
Solution
class Solution {public int numTeams(int[] rating) { int n = rating.length; int ans = 0; for (int i = 0; i < n; ++i) { int l = 0, r = 0; for (int j = 0; j < i; ++j) { if (rating[j] < rating[i]) { ++l; } } for (int j = i + 1; j < n; ++j) { if (rating[j] > rating[i]) { ++r; } } ans += l * r; ans += (i - l) * (n - i - 1 - r); } return ans; }}class Solution {public: int numTeams(vector<int> &rating) { int n = rating.size(); int ans = 0; for (int i = 0; i < n; ++i) { int l = 0, r = 0; for (int j = 0; j < i; ++j) { if (rating[j] < rating[i]) { ++l; } } for (int j = i + 1; j < n; ++j) { if (rating[j] > rating[i]) { ++r; } } ans += l * r; ans += (i - l) * (n - i - 1 - r); } return ans; }};class Solution: def numTeams(self, rating: List[int]) -> int: ans, n = 0, len(rating) for i, b in enumerate(rating): l = sum(a < b for a in rating[: i]) r = sum(c > b for c in rating[i + 1:]) ans += l * r ans += (i - l) * (n - i - 1 - r) return ansVideo walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.