1395. Count Number of Teams

MediumBinary Indexed TreeArrayDynamic Programming
Leetcode Link

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], and rating[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:

  1. We store all unique ratings in a sorted array to handle ratings[i] in a relative manner.
  2. We create two BITs, tree1 and tree2, for cumulative frequency calculations.
  3. Before we start iterating over each soldier, we update tree2 with their ratings to represent the total number of soldiers having a certain rating.
  4. For each soldier i in rating, 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 index x because we now have one more soldier with a rating up to x.
    • Query BIT tree1 to get the count of soldiers to the left of soldier i with a smaller rating (l), and similarly, query tree2 to get the count of soldiers to the right of i with a greater rating (r).
    • Calculate the number of valid increasing or decreasing triplets that soldier i can form as the center by multiplying l and r for the increasing sequence and (i-l) and (n-i-1-r) for the decreasing sequence, adding both results to the total count ans.
    • Update tree2 by subtracting 1 from index x because we have now counted soldier i and it should not be included in the calculation for future soldiers.

Finally, ans will contain the total number of valid teams that meet the problem's conditions.

Learn more about Dynamic Programming patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece:

What is the space complexity of the following code?

1int sum(int n) {
2  if (n <= 0) {
3    return 0;
4  }
5  return n + sum(n - 1);
6}

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:

  1. Initialization:

    • We first create a set from the rating array and then convert it into a sorted list called nums. This step effectively maps the actual ratings to indexes, where nums[0] would contain the smallest rating, and so on.
    • Two Binary Indexed Trees, tree1 and tree2, are initialized. tree1 is used to track the counts of soldiers when traversing from left to right, and tree2 would be used in the opposite direction.
  2. 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).
  3. Iterating Over Soldiers:

    • The main loop of the implementation goes over each soldier's rating. The variable n holds the total number of soldiers, and ans 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 the nums 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 index x-1.
      • Find the number of soldiers to the right of the current soldier with a higher rating (r by querying tree2 from index x+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) to ans. 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 index x 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 index x with -1 to exclude it from the right count for subsequent soldiers.

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.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

What is the space complexity of the following code?

1int sum(int n) {
2  if (n <= 0) {
3    return 0;
4  }
5  return n + sum(n - 1);
6}

Example Walkthrough

Let us take an example rating array: [2, 5, 3, 4, 1].

Following the solution approach:

  1. 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 and tree2.
  2. Populating tree2:

    • Before the iteration, we populate tree2 using the indices of nums. The frequencies of the ratings in tree2 will look like this: {1: 1, 2: 1, 3: 1, 4: 1, 5: 1} since each rating appears exactly once.
  3. Iterating Over Soldiers:

    • Start iterating over each soldier's rating at index i and ans is set to 0.
    • For i = 0 (rating 2):
      • Find the index x: 1 (nums[1] = 2).
      • Find l (left count) using tree1: query(0) = 0 (no soldiers to the left with lower ratings).
      • Find r (right count) using tree2: 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 to ans.
      • Update tree1 at index x: update(1, 1).
      • Update tree2 at index x: update(1, -1).
    • Continue for the rest of the indices:
      • i = 1 (rating 5): x = 4 ...
      • i = 2 (rating 3): x = 2 ...
      • i = 3 (rating 4): x = 3 ...
      • i = 4 (rating 1): x = 0 ...

    For the purpose of the example walkthrough, let's detail the step for i = 2 (rating 3):

    • Find the index x: 2 (nums[2] = 3).
    • Find l using tree1: query(1) = 1 (one soldier to the left with a lower rating 2).
    • Find r using tree2: query(5) - query(2) = 2 - 1 = 1 (one soldier to the right with a higher rating 4).
    • Add l * r + (2 - 1) * (4 - 2 - 1 - 1) = 1 * 1 + 1 * 0 = 1 to ans.
    • Update tree1 at index x with 1 and tree2 at index x with -1.

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.

Not Sure What to Study? Take the 2-min Quiz:

How does quick sort divide the problem into subproblems?

Python Solution

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

Java Solution

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

C++ Solution

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

Typescript Solution

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
Fast Track Your Learning with Our Quick Skills Quiz:

What is the space complexity of the following code?

1int sum(int n) {
2  if (n <= 0) {
3    return 0;
4  }
5  return n + sum(n - 1);
6}

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 for tree1 and tree2, where n is the length of the provided ratings array.
  • Sorting the distinct values in rating takes O(m log m) time, where m is the number of distinct ratings.
  • Updating and querying operations on a BIT take O(log m), where m is the number of nodes (or the length of the BIT array).
  • The for loop to update tree2 runs n times, making the complexity O(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 taking O(log m). This results into another O(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 and tree2 is O(m) each since each tree holds an array of size m + 1.
  • The space used for sorting to create the nums array is O(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.


Recommended Readings


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫