medium Count Number of Teams

Problem Statement

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

You need to form a team of 3 soldiers under these rules:

• Choose three soldiers with indices (i, j, k) and ratings (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 count of the valid teams that can be formed under these conditions.

Examples

Example 1:

• Input: rating = [3, 6, 2, 8, 1, 5]
• Expected Output: 3
• Justification: The valid teams are:
  • (3, 2, 1)
  • (3, 6, 8)
  • (6, 2, 1)

Example 2:

• Input: rating = [1, 5, 3, 2, 6, 4]
• Expected Output: 6
• Justification: The valid teams are:
  • (1, 3, 6)
  • (1, 2, 6)
  • (1, 2, 4)
  • (5, 3, 2)
  • (5, 3, 4)
  • (1, 5, 6)

Example 3:

• Input: rating = [10, 20, 30, 25, 15]
• Expected Output: 3
• Justification: The valid teams are:
  • (10, 20, 30)
  • (10, 20, 25)
  • (30, 25, 15)

Constraints:

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

Solution

To solve this problem, we use segment trees to efficiently count the number of valid teams. The idea is to use two segment trees: one for keeping track of the number of ratings to the left of the current soldier and another for ratings to the right. This allows us to efficiently query the number of ratings that are either less than or greater than the current rating in constant time.

The segment trees are updated as we iterate through the array. For each soldier, we count the number of valid teams by combining the counts from the left and right segment trees. By maintaining these counts, we can quickly determine the number of valid teams without having to repeatedly iterate over the entire array, making the solution more efficient.

Step-by-Step Algorithm

1. Initialize Variables:

   • Step 1.1: Initialize teamCount to 0. This variable will store the total number of valid teams.
   • Step 1.2: Calculate maxRating as the maximum rating in the ratings array plus 1. This helps in sizing the segment trees correctly to handle the range of rating values.
   • Step 1.3: Initialize leftSegmentTree and rightSegmentTree as integer arrays of size 4 * maxRating. These trees will keep track of the counts of ratings to the left and right of the current rating, respectively.

2. Build Initial Segment Trees:

   • Step 2.1: Iterate through the ratings array starting from the second element to the last element.
   • Step 2.2: For each rating in this iteration, update the rightSegmentTree by adding one to the count of the current rating using the update method. This step initializes the right segment tree with counts of all ratings except the first. This setup allows us to easily keep track of the counts of ratings to the right of the current rating as we iterate through the array.

3. Update Left Segment Tree with First Rating:

   • Step 3.1: Update the leftSegmentTree by adding one to the count of the first rating using the update method. This step sets up the left segment tree to start tracking the counts of ratings to the left of the current rating as we iterate through the array.

4. Iterate Through Ratings:

   • Step 4.1: Iterate through the ratings array starting from the second element to the last element. For each rating at index i, perform the following steps:

5. Count Teams with a Lower Rating on the Left and a Higher Rating on the Right:

   • Step 5.1: Use the sumRange method to query the leftSegmentTree for the number of ratings less than the current rating. This count is stored in leftLessCount.
   • Step 5.2: Use the sumRange method to query the rightSegmentTree for the number of ratings greater than the current rating. This count is stored in rightGreaterCount.
   • Step 5.3: Multiply leftLessCount and rightGreaterCount and add the result to teamCount. This step counts the valid teams where the current rating is in the middle, with a lower rating on the left and a higher rating on the right.

6. Count Teams with a Higher Rating on the Left and a Lower Rating on the Right:

   • Step 6.1: Use the sumRange method to query the leftSegmentTree for the number of ratings greater than the current rating. This count is stored in leftGreaterCount.
   • Step 6.2: Use the sumRange method to query the rightSegmentTree for the number of ratings less than the current rating. This count is stored in rightLessCount.
   • Step 6.3: Multiply leftGreaterCount and rightLessCount and add the result to teamCount. This step counts the valid teams where the current rating is in the middle, with a higher rating on the left and a lower rating on the right.

7. Update Segment Trees:

   • Step 7.1: Update the leftSegmentTree by adding one to the count of the current rating using the update method. This step updates the left segment tree to include the current rating, as it will be to the left of subsequent ratings.
   • Step 7.2: Update the rightSegmentTree by subtracting one from the count of the current rating using the update method. This step updates the right segment tree to exclude the current rating, as it has been processed and will no longer be to the right of any subsequent ratings.

8. Return the Total Number of Teams:

   • Step 8.1: Return the value of teamCount, which now contains the total number of valid teams.

Explanation of update() Method

• Step 1: If low equals high, it means we are at a leaf node in the segment tree. In this case, update the value at this node by adding value to segTree[index].
• Step 2: If low is not equal to high, calculate the middle point mid as (low + high) / 2.
• Step 3: If position lies between low and mid, recursively call update on the left child (2 * index + 1) with updated parameters.
• Step 4: If position lies between mid + 1 and high, recursively call update on the right child (2 * index + 2) with updated parameters.
• Step 5: After updating the appropriate child, update the current node value as the sum of its two children (segTree[index] = segTree[2 * index + 1] + segTree[2 * index + 2]).

Explanation of sumRange() Method

• Step 1: If low to high is completely within the left to right range, return segTree[index].
• Step 2: If there is no overlap between low to high and left to right, return 0.
• Step 3: If there is a partial overlap, calculate the middle point mid as (low + high) / 2.
• Step 4: Recursively query the left and right children (2 * index + 1 and 2 * index + 2) and sum their results to get the total range sum for the current node's range.

Algorithm Walkthrough

Input: [3, 6, 2, 8, 1, 5]

1. Initialize Variables:

• teamCount = 0
• maxRating = 9 (since max rating is 8 + 1)
• leftSegmentTree = new int[36] (4 * maxRating)
• rightSegmentTree = new int[36] (4 * maxRating)

2. Build Initial Right Segment Tree:

• For rating 6 (index 1):

  • Call update(0, 1, 0, 8, 6, rightSegmentTree)
    • low = 0, high = 8, mid = 4
    • 6 > 4, so call update(2, 1, 5, 8, 6, rightSegmentTree)
      • low = 5, high = 8, mid = 6
      • 6 <= 6, so call update(5, 1, 5, 6, 6, rightSegmentTree)
        • low = 5, high = 6, mid = 5
        • 6 > 5, so call update(12, 1, 6, 6, 6, rightSegmentTree)
          • low = 6, high = 6 (leaf node), update rightSegmentTree[12] += 1
      • Update rightSegmentTree[5] = rightSegmentTree[11] + rightSegmentTree[12] = 0 + 1 = 1
    • Update rightSegmentTree[2] = rightSegmentTree[5] + rightSegmentTree[6] = 1 + 0 = 1
  • Update rightSegmentTree[0] = rightSegmentTree[1] + rightSegmentTree[2] = 0 + 1 = 1

• For rating 2 (index 2):

  • Call update(0, 1, 0, 8, 2, rightSegmentTree)
    • low = 0, high = 8, mid = 4
    • 2 <= 4, so call update(1, 1, 0, 4, 2, rightSegmentTree)
      • low = 0, high = 4, mid = 2
      • 2 <= 2, so call update(3, 1, 0, 2, 2, rightSegmentTree)
        • low = 0, high = 2, mid = 1
        • 2 > 1, so call update(8, 1, 2, 2, 2, rightSegmentTree)
          • low = 2, high = 2 (leaf node), update rightSegmentTree[8] += 1
      • Update rightSegmentTree[3] = rightSegmentTree[7] + rightSegmentTree[8] = 0 + 1 = 1
    • Update rightSegmentTree[1] = rightSegmentTree[3] + rightSegmentTree[4] = 1 + 0 = 1
  • Update rightSegmentTree[0] = rightSegmentTree[1] + rightSegmentTree[2] = 1 + 1 = 2

• Similarly, process ratings 8, 1, and 5 to get the following rightSegmentTree:

  [
    5, 2, 3, 2, 0, 2, 1, 1, 1,
    0, 0, 1, 1, 0, 1, 0, 1, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0
  ]
  

3. Update Left Segment Tree with First Rating:

• For rating 3 (index 0):
  • Call update(0, 1, 0, 8, 3, leftSegmentTree)
    • low = 0, high = 8, mid = 4
    • 3 <= 4, so call update(1, 1, 0, 4, 3, leftSegmentTree)
      • low = 0, high = 4, mid = 2
      • 3 > 2, so call update(4, 1, 3, 4, 3, leftSegmentTree)
        • low = 3, high = 4, mid = 3
        • 3 <= 3, so call update(9, 1, 3, 3, 3, leftSegmentTree)
          • low = 3, high = 3 (leaf node), update leftSegmentTree[9] += 1
      • Update leftSegmentTree[4] = leftSegmentTree[9] + leftSegmentTree[10] = 1 + 0 = 1
    • Update leftSegmentTree[1] = leftSegmentTree[3] + leftSegmentTree[4] = 0 + 1 = 1
  • Update leftSegmentTree[0] = leftSegmentTree[1] + leftSegmentTree[2] = 1 + 0 = 1

  [
    1, 1, 0, 0, 1, 0, 0, 0, 0,
    1, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0
  ]
  

4. Iterate Through Ratings:

Step for i = 1 (Rating = 6):

Count teams with a lower rating on the left and a higher rating on the right:

• SumRange for leftLessCount:

  • Call sumRange(0, 5, 0, 8, 0, leftSegmentTree)
    • low = 0, high = 8, mid = 4
    • 5 <= 4: Call sumRange(0, 5, 0, 4, 1, leftSegmentTree)
      • low = 0, high = 4, mid = 2
      • 5 > 2: Call sumRange(0, 5, 3, 4, 4, leftSegmentTree)
        • low = 3, high = 4, mid = 3
        • 5 > 3: Call sumRange(0, 5, 4, 4, 10, leftSegmentTree)
          • low = 4, high = 4 (total overlap), return 1
      • leftLessCount = 1

• SumRange for rightGreaterCount:

  • Call sumRange(7, 8, 0, 8, 0, rightSegmentTree)
    • low = 0, high = 8, mid = 4
    • 7 > 4: Call sumRange(7, 8, 5, 8, 2, rightSegmentTree)
      • low = 5, high = 8, mid = 6
      • 7 > 6: Call sumRange(7, 8, 7, 8, 6, rightSegmentTree)
        • low = 7, high = 8, mid = 7
        • 7 <= 8: Call sumRange(7, 8, 8, 8, 14, rightSegmentTree)
          • low = 8, high = 8 (total overlap), return 1
      • rightGreaterCount = 1

• teamCount += 1 * 1 = 1

Count teams with a higher rating on the left and a lower rating on the right:

• SumRange for leftGreaterCount:

  • Call sumRange(7, 8, 0, 8, 0, leftSegmentTree)
    • low = 0, high = 8, mid = 4
    • 7 > 4: Call sumRange(7, 8, 5, 8, 2, leftSegmentTree)
      • low = 5, high = 8, mid = 6
      • 7 > 6: Call sumRange(7, 8, 7, 8, 6, leftSegmentTree)
        • low = 7, high = 8, mid = 7
        • 7 <= 8: Call sumRange(7, 8, 8, 8, 14, leftSegmentTree)
          • low = 8, high = 8 (total overlap), return 0
      • leftGreaterCount = 0

• SumRange for rightLessCount:

  • Call sumRange(0, 5, 0, 8, 0, rightSegmentTree)
    • low = 0, high = 8, mid = 4
    • 5 <= 4: Call sumRange(0, 5, 0, 4, 1, rightSegmentTree)
      • low = 0, high = 4, mid = 2
      • 5 > 2: Call sumRange(0, 5, 3, 4, 4, rightSegmentTree)
        • low = 3, high = 4, mid = 3
        • 5 > 3: Call sumRange(0, 5, 4, 4, 10, rightSegmentTree)
          • low = 4, high = 4 (total overlap), return 1
      • rightLessCount = 3

• teamCount = teamCount + 0 * 3 + 1 + 0 = 1

Update Segment Trees:

• Update leftSegmentTree at position 6
  • Call update(0, 1, 0, 8, 6, leftSegmentTree)
• Update rightSegmentTree at position 6
  • Call update(0, -1, 0, 8, 6, rightSegmentTree)

Step for i = 2 (Rating = 2):

• Count teams with a lower rating on the left and a higher rating on the right:

  • leftLessCount = sumRange(0, 1, 0, 8, 0, leftSegmentTree) = 0
  • rightGreaterCount = sumRange(3, 8, 0, 8, 0, rightSegmentTree) = 2
  • teamCount = teamCount + 0 * 2 = 1 + 0 = 1

• Count teams with a higher rating on the left and a lower rating on the right:

  • leftGreaterCount = sumRange(3, 8, 0, 8, 0, leftSegmentTree) = 2
  • rightLessCount = sumRange(0, 1, 0, 8, 0, rightSegmentTree) = 1
  • teamCount = teamCount + 2 * 1 = 1 + 2 = 3

Update Segment Trees:

• Update leftSegmentTree at position 2
  • Call update(0, 1, 0, 8, 2, leftSegmentTree)
• Update rightSegmentTree at position 2
  • Call update(0, -1, 0, 8, 2, rightSegmentTree)

Similarly process other array elements.

Final Output:

• teamCount = 3

Code

import java.util.Arrays;

class Solution {

  // Update the segment tree at a specific index
  public void update(
    int index,
    int value,
    int low,
    int high,
    int position,
    int[] segTree
  ) {
    if (low == high) {
      // Base case: Single element
      segTree[index] += value; // Update the value at the leaf node
    } else {
      int mid = (low + high) / 2; // Find the middle point
      // If position is in the left subtree
      if (position <= mid && position >= low) update(
        2 * index + 1,
        value,
        low,
        mid,
        position,
        segTree
      );
      else update(2 * index + 2, value, mid + 1, high, position, segTree);
      // Update the current node
      segTree[index] = segTree[2 * index + 1] + segTree[2 * index + 2];
    }
  }

  // Get the sum range in the segment tree
  private int sumRange(
    int left,
    int right,
    int low,
    int high,
    int index,
    int[] segTree
  ) {
    if (low >= left && high <= right) {
      // Total overlap
      return segTree[index];
    }
    if (right < low || high < left) {
      // No overlap
      return 0;
    }
    int mid = (low + high) / 2; // Partial overlap
    // Sum the ranges from both children
    return (
      sumRange(left, right, low, mid, 2 * index + 1, segTree) +
      sumRange(left, right, mid + 1, high, 2 * index + 2, segTree)
    );
  }

  // Main method to count the number of valid teams
  public int numTeams(int[] ratings) {
    int teamCount = 0; // Initialize team count
    int maxRating = Arrays.stream(ratings).max().getAsInt() + 1; // Find max rating
    int[] leftSegmentTree = new int[4 * maxRating]; // Segment tree for ratings to the left
    int[] rightSegmentTree = new int[4 * maxRating]; // Segment tree for ratings to the right

    // Update right segment tree with all ratings except the first
    for (int i = 1; i < ratings.length; i++) {
      update(0, 1, 0, maxRating - 1, ratings[i], rightSegmentTree);
    }
    // Update left segment tree with the first rating
    update(0, 1, 0, maxRating - 1, ratings[0], leftSegmentTree);

    // Iterate through ratings from the second element
    for (int i = 1; i < ratings.length; i++) {
      // Count teams with a lower rating on the left and a higher rating on the right
      int leftLessCount = sumRange(
        0,
        ratings[i] - 1,
        0,
        maxRating - 1,
        0,
        leftSegmentTree
      );
      int rightGreaterCount = sumRange(
        ratings[i] + 1,
        maxRating - 1,
        0,
        maxRating - 1,
        0,
        rightSegmentTree
      );
      teamCount += leftLessCount * rightGreaterCount;

      // Count teams with a higher rating on the left and a lower rating on the right
      int leftGreaterCount = sumRange(
        ratings[i] + 1,
        maxRating - 1,
        0,
        maxRating - 1,
        0,
        leftSegmentTree
      );
      int rightLessCount = sumRange(
        0,
        ratings[i] - 1,
        0,
        maxRating - 1,
        0,
        rightSegmentTree
      );
      teamCount += leftGreaterCount * rightLessCount;

      // Update the segment trees with the current rating
      update(0, 1, 0, maxRating - 1, ratings[i], leftSegmentTree);
      update(0, -1, 0, maxRating - 1, ratings[i], rightSegmentTree);
    }

    return teamCount; // Return the total number of valid teams
  }

  // Main method to test the solution
  public static void main(String[] args) {
    Solution solution = new Solution();

    // Test case 1
    int[] rating1 = { 3, 6, 2, 8, 1, 5 };
    int result1 = solution.numTeams(rating1);
    System.out.println("Expected Output: 3, Actual Output: " + result1);

    // Test case 2
    int[] rating2 = { 1, 5, 3, 2, 6, 4 };
    int result2 = solution.numTeams(rating2);
    System.out.println("Expected Output: 6, Actual Output: " + result2);

    // Test case 3
    int[] rating3 = { 10, 20, 30, 25, 15 };
    int result3 = solution.numTeams(rating3);
    System.out.println("Expected Output: 3, Actual Output: " + result3);
  }
}

Complexity Analysis

Time Complexity

• Building the Segment Tree: Each update operation on the segment tree takes time. Since we perform an update for each element in the array, the total time for building the segment tree is .
• Querying the Segment Tree: Each range sum query also takes time. We perform two range sum queries for each element in the array, so the total time for querying is .

Thus, the overall time complexity of the solution is .

Space Complexity

• Segment Trees: We use two segment trees, each of size 4 * maxRating. Since maxRating is the maximum value in the ratings array incremented by 1, the space complexity for the segment trees is .

Thus, the overall space complexity of the solution is .