Binary Indexed Tree
Segment Tree
Array
Dynamic Programming
Description
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.length
3 <= n <= 1000
1 <= rating[i] <= 105
All the integers in rating
are unique .
Solutions
Solution 1: Enumerate Middle Element
We can enumerate each element $rating[i]$ in the array $rating$ as the middle element, then count the number of elements $l$ that are smaller than it on the left, and the number of elements $r$ that are larger than it on the right. The number of combat units with this element as the middle element is $l \times r + (i - l) \times (n - i - 1 - r)$. We can add this to the answer.
The time complexity is $O(n^2)$, and the space complexity is $O(1)$. Where $n$ is the length of the array $rating$.
Python3 Java C++ Go TypeScript
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 ans
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22 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 ;
}
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23 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 ;
}
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19 func numTeams ( rating [] int ) ( ans int ) {
n := len ( rating )
for i , b := range rating {
l , r := 0 , 0
for _ , a := range rating [: i ] {
if a < b {
l ++
}
}
for _ , c := range rating [ i + 1 :] {
if c < b {
r ++
}
}
ans += l * r
ans += ( i - l ) * ( n - i - 1 - r )
}
return
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21 function numTeams ( rating : number []) : number {
let ans = 0 ;
const n = rating . length ;
for ( let i = 0 ; i < n ; ++ i ) {
let l = 0 ;
let r = 0 ;
for ( let j = 0 ; j < i ; ++ j ) {
if ( rating [ j ] < rating [ i ]) {
++ l ;
}
}
for ( let j = i + 1 ; j < n ; ++ j ) {
if ( rating [ j ] > rating [ i ]) {
++ r ;
}
}
ans += l * r ;
ans += ( i - l ) * ( n - i - 1 - r );
}
return ans ;
}
Solution 2: Binary Indexed Tree
We can use two binary indexed trees to maintain the number of elements $l$ that are smaller than each element on the left in the array $rating$, and the number of elements $r$ that are larger than it on the right. Then count the number of combat units with this element as the middle element as $l \times r + (i - l) \times (n - i - 1 - r)$, and add this to the answer.
The time complexity is $O(n \times \log n)$, and the space complexity is $O(n)$. Where $n$ is the length of the array $rating$.
Python3 Java C++ Go TypeScript
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38 class BinaryIndexedTree :
def __init__ ( self , n : int ):
self . n = n
self . c = [ 0 ] * ( n + 1 )
def update ( self , x : int , v : int ):
while x <= self . n :
self . c [ x ] += v
x += x & - x
def query ( self , x : int ) -> int :
s = 0
while x :
s += self . c [ x ]
x -= x & - x
return s
class Solution :
def numTeams ( self , rating : List [ int ]) -> int :
nums = sorted ( set ( rating ))
m = len ( nums )
tree1 = BinaryIndexedTree ( m )
tree2 = BinaryIndexedTree ( m )
for v in rating :
x = bisect_left ( nums , v ) + 1
tree2 . update ( x , 1 )
n = len ( rating )
ans = 0
for i , v in enumerate ( rating ):
x = bisect_left ( nums , v ) + 1
tree1 . update ( x , 1 )
tree2 . update ( x , - 1 )
l = tree1 . query ( x - 1 )
r = n - i - 1 - tree2 . query ( x )
ans += l * r
ans += ( i - l ) * ( n - i - 1 - r )
return ans
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70 class BinaryIndexedTree {
private int n ;
private int [] c ;
public BinaryIndexedTree ( int n ) {
this . n = n ;
this . c = new int [ n + 1 ] ;
}
public void update ( int x , int v ) {
while ( x <= n ) {
c [ x ] += v ;
x += x & - x ;
}
}
public int query ( int x ) {
int s = 0 ;
while ( x > 0 ) {
s += c [ x ] ;
x -= x & - x ;
}
return s ;
}
}
class Solution {
public int numTeams ( int [] rating ) {
int n = rating . length ;
int [] nums = rating . clone ();
Arrays . sort ( nums );
int m = 0 ;
for ( int i = 0 ; i < n ; ++ i ) {
if ( i == 0 || nums [ i ] != nums [ i - 1 ] ) {
nums [ m ++] = nums [ i ] ;
}
}
BinaryIndexedTree tree1 = new BinaryIndexedTree ( m );
BinaryIndexedTree tree2 = new BinaryIndexedTree ( m );
for ( int v : rating ) {
int x = search ( nums , v );
tree2 . update ( x , 1 );
}
int ans = 0 ;
for ( int i = 0 ; i < n ; ++ i ) {
int x = search ( nums , rating [ i ] );
tree1 . update ( x , 1 );
tree2 . update ( x , - 1 );
int l = tree1 . query ( x - 1 );
int r = n - i - 1 - tree2 . query ( x );
ans += l * r ;
ans += ( i - l ) * ( n - i - 1 - r );
}
return ans ;
}
private int search ( int [] nums , int x ) {
int l = 0 , r = nums . length ;
while ( l < r ) {
int mid = ( l + r ) >> 1 ;
if ( nums [ mid ] >= x ) {
r = mid ;
} else {
l = mid + 1 ;
}
}
return l + 1 ;
}
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54 class BinaryIndexedTree {
public :
BinaryIndexedTree ( int _n )
: n ( _n )
, c ( _n + 1 ) {}
void update ( int x , int delta ) {
while ( x <= n ) {
c [ x ] += delta ;
x += x & - x ;
}
}
int query ( int x ) {
int s = 0 ;
while ( x ) {
s += c [ x ];
x -= x & - x ;
}
return s ;
}
private :
int n ;
vector < int > c ;
};
class Solution {
public :
int numTeams ( vector < int >& rating ) {
vector < int > nums = rating ;
sort ( nums . begin (), nums . end ());
nums . erase ( unique ( nums . begin (), nums . end ()), nums . end ());
int m = nums . size ();
BinaryIndexedTree tree1 ( m );
BinaryIndexedTree tree2 ( m );
for ( int & v : rating ) {
int x = lower_bound ( nums . begin (), nums . end (), v ) - nums . begin () + 1 ;
tree2 . update ( x , 1 );
}
int ans = 0 ;
int n = rating . size ();
for ( int i = 0 ; i < n ; ++ i ) {
int x = lower_bound ( nums . begin (), nums . end (), rating [ i ]) - nums . begin () + 1 ;
tree1 . update ( x , 1 );
tree2 . update ( x , -1 );
int l = tree1 . query ( x - 1 );
int r = n - i - 1 - tree2 . query ( x );
ans += l * r ;
ans += ( i - l ) * ( n - i - 1 - r );
}
return ans ;
}
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55 type BinaryIndexedTree struct {
n int
c [] int
}
func newBinaryIndexedTree ( n int ) * BinaryIndexedTree {
c := make ([] int , n + 1 )
return & BinaryIndexedTree { n , c }
}
func ( this * BinaryIndexedTree ) update ( x , delta int ) {
for x <= this . n {
this . c [ x ] += delta
x += x & - x
}
}
func ( this * BinaryIndexedTree ) query ( x int ) int {
s := 0
for x > 0 {
s += this . c [ x ]
x -= x & - x
}
return s
}
func numTeams ( rating [] int ) ( ans int ) {
nums := make ([] int , len ( rating ))
copy ( nums , rating )
sort . Ints ( nums )
m := 0
for i , x := range nums {
if i == 0 || x != nums [ i - 1 ] {
nums [ m ] = x
m ++
}
}
nums = nums [: m ]
tree1 := newBinaryIndexedTree ( m )
tree2 := newBinaryIndexedTree ( m )
for _ , x := range rating {
tree2 . update ( sort . SearchInts ( nums , x ) + 1 , 1 )
}
n := len ( rating )
for i , v := range rating {
x := sort . SearchInts ( nums , v ) + 1
tree1 . update ( x , 1 )
tree2 . update ( x , - 1 )
l := tree1 . query ( x - 1 )
r := n - i - 1 - tree2 . query ( x )
ans += l * r
ans += ( i - l ) * ( n - i - 1 - r )
}
return
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67 class BinaryIndexedTree {
private n : number ;
private c : number [];
constructor ( n : number ) {
this . n = n ;
this . c = new Array ( n + 1 ). fill ( 0 );
}
public update ( x : number , v : number ) : void {
while ( x <= this . n ) {
this . c [ x ] += v ;
x += x & - x ;
}
}
public query ( x : number ) : number {
let s = 0 ;
while ( x > 0 ) {
s += this . c [ x ];
x -= x & - x ;
}
return s ;
}
}
function numTeams ( rating : number []) : number {
let nums = [... rating ];
nums . sort (( a , b ) => a - b );
const n = rating . length ;
let m = 0 ;
for ( let i = 0 ; i < n ; ++ i ) {
if ( i === 0 || nums [ i ] !== nums [ i - 1 ]) {
nums [ m ++ ] = nums [ i ];
}
}
nums = nums . slice ( 0 , m );
const search = ( x : number ) : number => {
let l = 0 ;
let r = m ;
while ( l < r ) {
const mid = ( l + r ) >> 1 ;
if ( nums [ mid ] >= x ) {
r = mid ;
} else {
l = mid + 1 ;
}
}
return l + 1 ;
};
let ans = 0 ;
const tree1 = new BinaryIndexedTree ( m );
const tree2 = new BinaryIndexedTree ( m );
for ( const x of rating ) {
tree2 . update ( search ( x ), 1 );
}
for ( let i = 0 ; i < n ; ++ i ) {
const x = search ( rating [ i ]);
tree1 . update ( x , 1 );
tree2 . update ( x , - 1 );
const l = tree1 . query ( x - 1 );
const r = n - i - 1 - tree2 . query ( x );
ans += l * r ;
ans += ( i - l ) * ( n - i - 1 - r );
}
return ans ;
}
Solution 3: Recursion + Memoization
TypeScript JavaScript
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41 function numTeams ( rating : number []) : number {
const n = rating . length ;
const f : Record < Type , number [][] > = {
asc : Array.from ({ length : n }, () => Array ( 3 ). fill ( - 1 )),
desc : Array.from ({ length : n }, () => Array ( 3 ). fill ( - 1 )),
};
const fn = ( i : number , available : number , type : Type ) => {
if ( ! available ) {
return 1 ;
}
if ( f [ type ][ i ][ available ] !== - 1 ) {
return f [ type ][ i ][ available ];
}
let ans = 0 ;
for ( let j = i + 1 ; j < n ; j ++ ) {
if ( rating [ j ] > rating [ i ]) {
if ( type === 'asc' ) {
ans += fn ( j , available - 1 , 'asc' );
}
} else {
if ( type === 'desc' ) {
ans += fn ( j , available - 1 , 'desc' );
}
}
}
f [ type ][ i ][ available ] = ans ;
return ans ;
};
let ans = 0 ;
for ( let i = 0 ; i < n ; i ++ ) {
ans += fn ( i , 2 , 'asc' ) + fn ( i , 2 , 'desc' );
}
return ans ;
}
type Type = 'asc' | 'desc' ;
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43 /**
* @param {number[]} rating
* @return {number}
*/
var numTeams = function ( rating ) {
const n = rating . length ;
const f = {
asc : Array . from ({ length : n }, () => Array ( 3 ). fill ( - 1 )),
desc : Array . from ({ length : n }, () => Array ( 3 ). fill ( - 1 )),
};
const fn = ( i , available , type ) => {
if ( ! available ) {
return 1 ;
}
if ( f [ type ][ i ][ available ] !== - 1 ) {
return f [ type ][ i ][ available ];
}
let ans = 0 ;
for ( let j = i + 1 ; j < n ; j ++ ) {
if ( rating [ j ] > rating [ i ]) {
if ( type === 'asc' ) {
ans += fn ( j , available - 1 , 'asc' );
}
} else {
if ( type === 'desc' ) {
ans += fn ( j , available - 1 , 'desc' );
}
}
}
f [ type ][ i ][ available ] = ans ;
return ans ;
};
let ans = 0 ;
for ( let i = 0 ; i < n ; i ++ ) {
ans += fn ( i , 2 , 'asc' ) + fn ( i , 2 , 'desc' );
}
return ans ;
};
GitHub