Taro Logo

Count Number of Teams

Medium
Goldman Sachs logo
Goldman Sachs
0 views
Topics:
Arrays

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

Solution


Clarifying Questions

When you get asked this question in a real-life environment, it will often be ambiguous (especially at FAANG). Make sure to ask these questions in that case:

  1. What is the maximum possible length of the `rating` array?
  2. Can the integers in the `rating` array be negative, zero, or only positive?
  3. Are duplicate values allowed in the `rating` array, and if so, how should they be handled?
  4. If no valid teams can be formed, what value should I return?
  5. Is the range of values within the rating array going to be bounded to any particular range such as within the range of a 32-bit integer?

Brute Force Solution

Approach

The brute force strategy for counting teams is about checking every possible combination of soldiers. We explore all groups of three soldiers to see if they meet the team criteria by comparing their ratings.

Here's how the algorithm would work step-by-step:

  1. Consider each soldier as the potential first member of a team.
  2. For that first soldier, look at every possible pair of other soldiers that could complete the team.
  3. For each possible team of three soldiers, check if their ratings are in strictly increasing order.
  4. If not, check if their ratings are in strictly decreasing order.
  5. If either of those conditions is true, then that group of three soldiers forms a valid team, so count it.
  6. Repeat this process, trying every possible first soldier and every possible pair of teammates.
  7. Once you have checked all possible combinations, the total number of valid teams is your answer.

Code Implementation

def count_number_of_teams(ratings):
    number_of_teams = 0

    for first_soldier_index in range(len(ratings)):
        for second_soldier_index in range(first_soldier_index + 1, len(ratings)):
            for third_soldier_index in range(second_soldier_index + 1, len(ratings)):
                # We now have a possible team of three soldiers

                first_soldier_rating = ratings[first_soldier_index]
                second_soldier_rating = ratings[second_soldier_index]
                third_soldier_rating = ratings[third_soldier_index]

                # Check for increasing order
                if first_soldier_rating < second_soldier_rating < third_soldier_rating:
                    number_of_teams += 1

                # Check for decreasing order
                elif first_soldier_rating > second_soldier_rating > third_soldier_rating:
                    number_of_teams += 1

    return number_of_teams

Big(O) Analysis

Time Complexity
O(n^3)The algorithm iterates through each soldier as a potential starting member of a team, resulting in an outer loop that runs 'n' times, where 'n' is the number of soldiers. For each soldier in the outer loop, two more nested loops iterate through the remaining soldiers to form pairs. Therefore, for each starting soldier, we potentially consider every possible pair of the other soldiers, which leads to n*(n-1)*(n-2) operations, approximated as n*n*n. Simplifying, the time complexity is O(n^3).
Space Complexity
O(1)The provided brute force approach iterates through all possible combinations of three soldiers without using any auxiliary data structures. It only uses a few integer variables to store the indices of the soldiers being considered and a counter for the number of teams. The space required for these variables remains constant regardless of the number of soldiers, N. Therefore, the space complexity is O(1).

Optimal Solution

Approach

To count the number of teams, we'll consider each soldier as the middle person on a potential team. For each soldier, we'll count how many soldiers to their left are smaller and how many are larger, and similarly count the smaller and larger soldiers to their right. The number of teams that soldier can be a part of is the product of those counts.

Here's how the algorithm would work step-by-step:

  1. Go through each soldier in the line, one at a time.
  2. For the current soldier, count how many soldiers to their left are shorter than them.
  3. Still for the current soldier, count how many soldiers to their left are taller than them.
  4. Now, count how many soldiers to their right are shorter than the current soldier.
  5. Lastly, count how many soldiers to their right are taller than the current soldier.
  6. For each soldier, multiply the number of shorter soldiers on the left by the number of taller soldiers on the right. Also, multiply the number of taller soldiers on the left by the number of shorter soldiers on the right.
  7. Add up all of those products (calculated in the previous step) for each soldier. That final sum is the total number of possible teams.

Code Implementation

def count_number_of_teams(ratings):
    number_of_soldiers = len(ratings)
    team_count = 0

    for middle_soldier in range(1, number_of_soldiers - 1):
        smaller_on_left = 0
        larger_on_left = 0
        smaller_on_right = 0
        larger_on_right = 0

        # Count soldiers to the left.
        for left_soldier in range(middle_soldier):
            if ratings[left_soldier] < ratings[middle_soldier]:
                smaller_on_left += 1
            elif ratings[left_soldier] > ratings[middle_soldier]:
                larger_on_left += 1

        # Count soldiers to the right.
        for right_soldier in range(middle_soldier + 1, number_of_soldiers):
            if ratings[right_soldier] < ratings[middle_soldier]:
                smaller_on_right += 1
            elif ratings[right_soldier] > ratings[middle_soldier]:
                larger_on_right += 1

        # Teams can be formed in two ways.
        team_count += smaller_on_left * larger_on_right

        team_count += larger_on_left * smaller_on_right

    return team_count

Big(O) Analysis

Time Complexity
O(n²)The algorithm iterates through each soldier in the rating array, which has n elements. For each soldier, it iterates through the soldiers to their left and to their right to count taller and shorter soldiers. Thus, for each of the n soldiers, it performs up to n-1 comparisons. This results in a nested loop structure that performs approximately n * (n-1) operations, which simplifies to O(n²).
Space Complexity
O(1)The provided plain English explanation describes an iterative process that primarily involves counting taller and shorter soldiers to the left and right of each soldier. It does not mention creating any additional data structures like arrays, hash maps, or lists to store intermediate results. The algorithm only needs to store a fixed number of counters (smaller left, larger left, smaller right, larger right) for each soldier, and these counters do not scale with the input size N (the number of soldiers). Therefore, the auxiliary space complexity is constant.

Edge Cases

CaseHow to Handle
Empty or null rating arrayReturn 0 immediately as no team can be formed.
Rating array with fewer than 3 elementsReturn 0 immediately as a team requires 3 soldiers.
Rating array with all identical valuesThe result should be 0 since there cannot be a strictly increasing or decreasing sequence.
Rating array with a sorted increasing sequenceTest case will return (n choose 3) where n is the array length, which can be large.
Rating array with a sorted decreasing sequenceTest case will return (n choose 3) where n is the array length, which can be large.
Rating array with large numbers that might cause integer overflow if multipliedEnsure calculations are performed with a data type that prevents overflow (e.g., long).
Rating array with duplicate valuesThe comparisons should still function correctly, counting valid teams even with duplicate skill levels at different indices.
Maximum size rating array to test performanceVerify the algorithm scales efficiently (e.g., O(n^2) or better) to avoid exceeding time limits.