Skip to content

1395. Count Number of Teams

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\).

1
2
3
4
5
6
7
8
9
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\).

 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

 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;
};

Comments