1395. Count Number of Teams
Problem Description
The problem presented deals with forming teams of 3 soldiers from a line-up where each soldier has a unique rating value. The goal is to form a team following certain rules. A team is considered valid if one of these two conditions are satisfied:
- the ratings of the three soldiers (let's call them
rating[i]
,rating[j]
, andrating[k]
) are in increasing order (rating[i] < rating[j] < rating[k]
) - or the ratings are in decreasing order (
rating[i] > rating[j] > rating[k]
).
The key constraint is that we must choose three soldiers with indices i
, j
, k
(0 ≤ i < j < k < n
) to form a valid team. The question asks for the total number of valid teams that we can create, given that soldiers can be part of multiple teams.
Intuition
The solution uses a Binary Indexed Tree (BIT), also known as a Fenwick Tree, to efficiently calculate the cumulative frequencies (or, in this context, the number of soldiers with certain ratings) and modifies them over time.
To solve this problem, it becomes necessary to consider each valid triplet (i, j, k)
and check if they form an increasing or decreasing sequence. However, doing so with a brute force approach would take O(n^3)
time, which is too slow for large inputs.
Instead of checking every possible triplet, we leverage the BIT to store and quickly compute the number of ratings that are less than or greater than a given rating (while maintaining the ordering constraint). This is helpful because, for a given soldier j
, we are interested in how many soldiers to the left have a lower rating (for an increasing sequence) or a higher rating (for a decreasing sequence), and similarly, how many soldiers to the right have a higher rating (increasing sequence) or a lower rating (decreasing sequence).
Here’s the intuition behind our solution steps:
- We store all unique ratings in a sorted array to handle
ratings[i]
in a relative manner. - We create two BITs,
tree1
andtree2
, for cumulative frequency calculations. - Before we start iterating over each soldier, we update
tree2
with their ratings to represent the total number of soldiers having a certain rating. - For each soldier
i
inrating
, we do the following:- Using binary search, find the index
x
of the current rating in the sorted array of ratings. - Update BIT
tree1
by adding 1 to indexx
because we now have one more soldier with a rating up tox
. - Query BIT
tree1
to get the count of soldiers to the left of soldieri
with a smaller rating (l
), and similarly, querytree2
to get the count of soldiers to the right ofi
with a greater rating (r
). - Calculate the number of valid increasing or decreasing triplets that soldier
i
can form as the center by multiplyingl
andr
for the increasing sequence and(i-l)
and(n-i-1-r)
for the decreasing sequence, adding both results to the total countans
. - Update
tree2
by subtracting 1 from indexx
because we have now counted soldieri
and it should not be included in the calculation for future soldiers.
- Using binary search, find the index
Finally, ans
will contain the total number of valid teams that meet the problem's conditions.
Learn more about Dynamic Programming patterns.
Solution Approach
The implementation of the solution involves the usage of Binary Indexed Trees (BITs) to perform efficient range queries and updates. The Binary Indexed Tree is a data structure that allows getting the prefix sum up to a certain index and updating the frequencies, both operations in O(log n)
time. This feature is particularly useful in this problem when we need to calculate the count of soldiers with ratings less than or greater than the current rating quickly.
Here's a step-by-step breakdown of the solution's implementation:
-
Initialization:
- We first create a set from the
rating
array and then convert it into a sorted list callednums
. This step effectively maps the actual ratings to indexes, wherenums[0]
would contain the smallest rating, and so on. - Two Binary Indexed Trees,
tree1
andtree2
, are initialized.tree1
is used to track the counts of soldiers when traversing from left to right, andtree2
would be used in the opposite direction.
- We first create a set from the
-
Populating
tree2
:- Before iterating over the soldiers, we populate
tree2
with the frequency of each soldier's rating. This preparation step will help us later when we want to count the number of soldiers on the right of a particular soldier with a certain condition (rating higher or lower).
- Before iterating over the soldiers, we populate
-
Iterating Over Soldiers:
- The main loop of the implementation goes over each soldier's rating. The variable
n
holds the total number of soldiers, andans
is initialized to 0, to accumulate the number of valid teams. - For each soldier with rating
v
, we perform the following actions:- Find the index
x
representing the soldier's rating in thenums
array using binary search. - Use
tree1
to find the number of soldiers to the left of the current soldier with a lower rating (l
) by querying it up to indexx-1
. - Find the number of soldiers to the right of the current soldier with a higher rating (
r
by queryingtree2
from indexx+1
to the end and subtracting it from the total remaining soldiers to the right. - Add the products
l * r
and(i - l) * (n - i - 1 - r)
toans
. The first product corresponds to the number of increasing sequences that can be formed, and the second is for the decreasing sequences. - Update
tree1
by adding 1 to indexx
to include the current soldier's rating in the left count for subsequent soldiers. - Decrement the frequency of the current soldier's rating in
tree2
by updating indexx
with-1
to exclude it from the right count for subsequent soldiers.
- Find the index
- The main loop of the implementation goes over each soldier's rating. The variable
This approach efficiently computes the valid teams by focusing on each soldier and considering them as the middle soldier of a team, checking how many teams they can form with soldiers to their left and right.
After iterating through all the soldiers, the variable ans
holds the final count of all valid teams, according to the problem's conditions.
Ready to land your dream job?
Unlock your dream job with a 2-minute evaluator for a personalized learning plan!
Start EvaluatorExample Walkthrough
Let us take an example rating
array: [2, 5, 3, 4, 1]
.
Following the solution approach:
-
Initialization:
- We first create a set from the
rating
array:{1, 2, 3, 4, 5}
. - Convert this set into a sorted list
nums
:[1, 2, 3, 4, 5]
. - Initialize two Binary Indexed Trees,
tree1
andtree2
.
- We first create a set from the
-
Populating
tree2
:- Before the iteration, we populate
tree2
using the indices ofnums
. The frequencies of the ratings intree2
will look like this:{1: 1, 2: 1, 3: 1, 4: 1, 5: 1}
since each rating appears exactly once.
- Before the iteration, we populate
-
Iterating Over Soldiers:
- Start iterating over each soldier's rating at index
i
andans
is set to 0. - For
i = 0
(rating2
):- Find the index
x
:1
(nums[1]
= 2). - Find
l
(left count) usingtree1
:query(0) = 0
(no soldiers to the left with lower ratings). - Find
r
(right count) usingtree2
:query(5) - query(1) = 4 - 1 = 3
(three soldiers to the right with higher ratings). - Add
l * r + (0 - l) * (4 - 0 - 1 - r) = 0 * 3 + 0 * 0 = 0
toans
. - Update
tree1
at indexx
:update(1, 1)
. - Update
tree2
at indexx
:update(1, -1)
.
- Find the index
- Continue for the rest of the indices:
i = 1
(rating5
):x = 4
...i = 2
(rating3
):x = 2
...i = 3
(rating4
):x = 3
...i = 4
(rating1
):x = 0
...
For the purpose of the example walkthrough, let's detail the step for
i = 2
(rating3
):- Find the index
x
:2
(nums[2]
= 3). - Find
l
usingtree1
:query(1) = 1
(one soldier to the left with a lower rating2
). - Find
r
usingtree2
:query(5) - query(2) = 2 - 1 = 1
(one soldier to the right with a higher rating4
). - Add
l * r + (2 - 1) * (4 - 2 - 1 - 1) = 1 * 1 + 1 * 0 = 1
toans
. - Update
tree1
at indexx
with 1 andtree2
at indexx
with -1.
- Start iterating over each soldier's rating at index
At the end of iteration, we add up all the contributions to ans
from each i
, which gives us the total count of valid teams.
In our example, the final ans
would be the total number of valid sequences that can be formed using the ratings provided, following the rules described in the problem.
Solution Implementation
1from typing import List
2from bisect import bisect_left
3
4# Binary Indexed Tree (Fenwick Tree) to perform efficient update and prefix sum queries
5class BinaryIndexedTree:
6 def __init__(self, size: int):
7 # Initialize the size and an internal list with an extra space (index 0 is unused)
8 self.size = size
9 self.tree_array = [0] * (size + 1)
10
11 # Adds a value v to the element at index x
12 def update(self, index: int, value: int):
13 while index <= self.size:
14 self.tree_array[index] += value
15 index += index & -index
16
17 # Calculates the prefix sum up to and including index x
18 def query(self, index: int) -> int:
19 sum = 0
20 while index:
21 sum += self.tree_array[index]
22 index -= index & -index
23 return sum
24
25# Solution class to identify the number of teams that satisfy the contest requirements
26class Solution:
27 def numTeams(self, rating: List[int]) -> int:
28 # Sort and remove duplicates to 'compress' the values and map them
29 unique_sorted_ratings = sorted(set(rating))
30 mapping_size = len(unique_sorted_ratings)
31
32 # Two trees to calculate the number of soldiers less and greater than a given rating
33 tree_smaller = BinaryIndexedTree(mapping_size)
34 tree_greater = BinaryIndexedTree(mapping_size)
35
36 # Initialize tree_greater with the count of each rating
37 for value in rating:
38 index = bisect_left(unique_sorted_ratings, value) + 1
39 tree_greater.update(index, 1)
40
41 total_soldiers = len(rating)
42 teams_count = 0
43 # Iterate through each soldier's rating
44 for i, value in enumerate(rating):
45 index = bisect_left(unique_sorted_ratings, value) + 1
46 # Update the trees as we move the 'split' line
47 tree_smaller.update(index, 1)
48 tree_greater.update(index, -1)
49
50 # Count soldiers smaller and greater than the current value to the left and right
51 left_smaller = tree_smaller.query(index - 1)
52 right_greater = total_soldiers - i - 1 - tree_greater.query(index)
53
54 # Accumulate the number of teams possible with the current soldier
55 teams_count += left_smaller * right_greater
56 teams_count += (i - left_smaller) * (total_soldiers - i - 1 - right_greater)
57
58 # Return the total count of possible teams
59 return teams_count
60
61# Example usage of the code, you can remove or comment out this part
62# (Note: The list 'rating' must be provided in the real LeetCode environment)
63# solution = Solution()
64# result = solution.numTeams(rating)
65# print(result)
66
1import java.util.Arrays;
2
3class BinaryIndexedTree {
4 // Size of the array
5 private int size;
6 // The underlying data structure for Binary Indexed Tree
7 private int[] tree;
8
9 // Constructor to initialize the tree with a given size n
10 public BinaryIndexedTree(int n) {
11 this.size = n;
12 this.tree = new int[n + 1];
13 }
14
15 // Method to update the Binary Indexed Tree with a value v at index x
16 public void update(int index, int value) {
17 while (index <= size) {
18 tree[index] += value;
19 index += index & -index;
20 }
21 }
22
23 // Method to query the prefix sum up to index x
24 public int query(int index) {
25 int sum = 0;
26 while (index > 0) {
27 sum += tree[index];
28 index -= index & -index;
29 }
30 return sum;
31 }
32}
33
34class Solution {
35 // Method to find the number of teams that satisfy the rating conditions
36 public int numTeams(int[] rating) {
37 int n = rating.length;
38 int[] sortedRating = rating.clone();
39 Arrays.sort(sortedRating);
40 int uniqueElements = 0;
41
42 // Reduce the sorted array to only unique elements
43 for (int i = 0; i < n; ++i) {
44 if (i == 0 || sortedRating[i] != sortedRating[i - 1]) {
45 sortedRating[uniqueElements++] = sortedRating[i];
46 }
47 }
48
49 // Initialize two Binary Indexed Trees
50 BinaryIndexedTree ascTree = new BinaryIndexedTree(uniqueElements);
51 BinaryIndexedTree descTree = new BinaryIndexedTree(uniqueElements);
52
53 // Populate descTree with the indices of rating values in sortedRating
54 for (int value : rating) {
55 int index = binarySearch(sortedRating, value);
56 descTree.update(index, 1);
57 }
58
59 int totalTeams = 0; // Stores the total number of valid teams
60 // Loop through each rating as the middle element of the team
61 for (int i = 0; i < n; ++i) {
62 int index = binarySearch(sortedRating, rating[i]);
63 ascTree.update(index, 1); // Update ascTree for current element
64 descTree.update(index, -1); // Decrement count in descTree as it's now used
65
66 // Calculate the number of lower elements before and higher elements after the current element
67 int leftLessCount = ascTree.query(index - 1);
68 int rightGreaterCount = n - i - 1 - descTree.query(index);
69
70 // Add to the total number of teams possible with current as middle rating
71 totalTeams += leftLessCount * rightGreaterCount;
72 // Also add the teams with current as the greatest rating
73 totalTeams += (i - leftLessCount) * (n - i - 1 - rightGreaterCount);
74 }
75
76 // Return the total number of possible teams
77 return totalTeams;
78 }
79
80 // Helper method to perform a binary search to find the index of x in nums
81 private int binarySearch(int[] nums, int target) {
82 int left = 0, right = nums.length;
83 while (left < right) {
84 int mid = (left + right) >> 1;
85 if (nums[mid] >= target) {
86 right = mid;
87 } else {
88 left = mid + 1;
89 }
90 }
91 return left + 1; // Addition of 1 to match the 1-indexed nature of BIT
92 }
93}
94
1#include <vector>
2#include <algorithm>
3
4using std::vector;
5using std::sort;
6using std::unique;
7using std::lower_bound;
8
9class BinaryIndexedTree {
10public:
11 // Constructor: initializes the Binary Indexed Tree with a given size.
12 explicit BinaryIndexedTree(int size)
13 : size_(size), tree_(size + 1, 0) {}
14
15 // Updates the Binary Indexed Tree with a given value at a specific position.
16 void update(int index, int delta) {
17 while (index <= size_) {
18 tree_[index] += delta;
19 index += index & -index; // Move to the next index to be updated.
20 }
21 }
22
23 // Queries the sum of values up to a given position.
24 int query(int index) {
25 int sum = 0;
26 while (index > 0) {
27 sum += tree_[index];
28 index -= index & -index; // Move to the parent index.
29 }
30 return sum;
31 }
32
33private:
34 int size_; // The size of the original input array.
35 vector<int> tree_; // The underlying tree represented as a vector.
36};
37
38class Solution {
39public:
40 // Counts the number of teams that satisfy the rating condition.
41 int numTeams(vector<int>& rating) {
42 // Create and sort a unique copy of the ratings to map values to indices.
43 vector<int> sortedRating = rating;
44 sort(sortedRating.begin(), sortedRating.end());
45 sortedRating.erase(unique(sortedRating.begin(), sortedRating.end()), sortedRating.end());
46
47 int uniqueSize = sortedRating.size();
48
49 // Initialize two Binary Indexed Trees.
50 BinaryIndexedTree treeAscending(uniqueSize);
51 BinaryIndexedTree treeDescending(uniqueSize);
52
53 // Fill in the treeDescending with initial counts.
54 for (int v : rating) {
55 int rank = lower_bound(sortedRating.begin(), sortedRating.end(), v) - sortedRating.begin() + 1;
56 treeDescending.update(rank, 1);
57 }
58
59 int result = 0;
60 int n = rating.size();
61 for (int i = 0; i < n; ++i) {
62 // Find the rank of the current rating in the unique list.
63 int rank = lower_bound(sortedRating.begin(), sortedRating.end(), rating[i]) - sortedRating.begin() + 1;
64
65 // Update Binary Indexed Trees.
66 treeAscending.update(rank, 1);
67 treeDescending.update(rank, -1);
68
69 // Query the counts for constructing teams.
70 int lessThanCurrent = treeAscending.query(rank - 1);
71 int greaterThanCurrent = n - i - 1 - treeDescending.query(rank);
72
73 // Calculate the number of valid teams.
74 result += lessThanCurrent * greaterThanCurrent; // Ascending on left, descending on right.
75 result += (i - lessThanCurrent) * (n - i - 1 - greaterThanCurrent); // Descending on left, ascending on right.
76 }
77
78 return result;
79 }
80};
81
1let n: number;
2let c: number[];
3
4// Initialize the Binary Indexed Tree with n elements.
5function init(nParam: number): void {
6 n = nParam;
7 c = new Array(n + 1).fill(0);
8}
9
10// Update the Binary Indexed Tree at index x with value v.
11function update(x: number, v: number): void {
12 while (x <= n) {
13 c[x] += v;
14 x += x & -x;
15 }
16}
17
18// Query the cumulative frequency up to and including index x.
19function query(x: number): number {
20 let sum = 0;
21 while (x > 0) {
22 sum += c[x];
23 x -= x & -x;
24 }
25 return sum;
26}
27
28// Calculate the number of valid teams that can be formed from the input array of ratings.
29function numTeams(rating: number[]): number {
30 let nums = [...rating];
31 nums.sort((a, b) => a - b);
32 const ratingLength = rating.length;
33 let uniqueCount = 0;
34 for (let i = 0; i < ratingLength; ++i) {
35 if (i === 0 || nums[i] !== nums[i - 1]) {
36 nums[uniqueCount++] = nums[i];
37 }
38 }
39 nums = nums.slice(0, uniqueCount);
40
41 // Binary search to find the index of x in nums.
42 const search = (x: number): number => {
43 let left = 0;
44 let right = uniqueCount;
45 while (left < right) {
46 const mid = (left + right) >> 1;
47 if (nums[mid] >= x) {
48 right = mid;
49 } else {
50 left = mid + 1;
51 }
52 }
53 return left + 1;
54 };
55
56 let totalTeams = 0;
57 init(uniqueCount); // Initialize Binary Indexed Trees
58 const tree1 = [...Array(uniqueCount + 1).fill(0)];
59 const tree2 = [...tree1];
60
61 for (const x of rating) {
62 update(search(x), 1);
63 }
64
65 for (let i = 0; i < ratingLength; ++i) {
66 const x = search(rating[i]);
67 update(x, 1);
68 const leftCount = query(x - 1);
69 const rightCount = ratingLength - i - 1 - query(x);
70 totalTeams += leftCount * rightCount;
71 totalTeams += (i - leftCount) * (ratingLength - i - 1 - rightCount);
72 }
73
74 return totalTeams;
75}
76
Time and Space Complexity
The code provided uses a Binary Indexed Tree (BIT) structure to solve the problem of counting the number of teams that can be formed according to the given rating constraints. Let's analyze the time complexity and space complexity of the code:
Time Complexity
- The initialization of the BinaryIndexedTree takes
O(n)
time, both fortree1
andtree2
, wheren
is the length of the provided ratings array. - Sorting the distinct values in
rating
takesO(m log m)
time, wherem
is the number of distinct ratings. - Updating and querying operations on a BIT take
O(log m)
, wherem
is the number of nodes (or the length of the BIT array). - The for loop to update
tree2
runsn
times, making the complexityO(n log m)
. - The main for loop also runs
n
times, with each iteration performing a constant number of updates and queries on the BIT, each takingO(log m)
. This results into anotherO(n log m)
.
The overall time complexity can be calculated by adding all the individual complexities:
O(n) + O(m log m) + O(n log m) + O(n log m) = O((2n + m) log m)
.
With the constraint that m <= n
, the expression simplifies to O(n log n)
.
Space Complexity
- The space complexity of the BIT
tree1
andtree2
isO(m)
each since each tree holds an array of sizem + 1
. - The space used for sorting to create the
nums
array isO(m)
. - Auxiliary space used for variables is
O(1)
.
Therefore, the total space complexity is primarily determined by the BIT arrays and the sorted list of unique elements. Hence, the final space complexity is O(m) + O(m) + O(m) = O(3m)
, which simplifies to O(m)
. Since m <= n
, this can also be simplified to O(n)
.
Both complexities are dependent on the size of the rating
input array and the number of distinct ratings within it.
Learn more about how to find time and space complexity quickly using problem constraints.
Which algorithm should you use to find a node that is close to the root of the tree?
Recommended Readings
What is Dynamic Programming Prerequisite DFS problems dfs_intro Backtracking problems backtracking Memoization problems memoization_intro Pruning problems backtracking_pruning Dynamic programming is an algorithmic optimization technique that breaks down a complicated problem into smaller overlapping sub problems in a recursive manner and uses solutions to the sub problems to construct a solution
LeetCode Patterns Your Personal Dijkstra's Algorithm to Landing Your Dream Job The goal of AlgoMonster is to help you get a job in the shortest amount of time possible in a data driven way We compiled datasets of tech interview problems and broke them down by patterns This way we
Recursion Recursion is one of the most important concepts in computer science Simply speaking recursion is the process of a function calling itself Using a real life analogy imagine a scenario where you invite your friends to lunch https algomonster s3 us east 2 amazonaws com recursion jpg You first
Want a Structured Path to Master System Design Too? Don’t Miss This!