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

Thinking

Count strictly increasing or decreasing triples. \(n \le 1000\) makes a triple loop cubic. Fix the middle rating \(b\): \(l\) smaller values on the left times \(r\) larger on the right are the increasing teams; the complementary counts are the decreasing ones. Two linear scans per midpoint suffice.

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

Thinking

Scanning both sides per midpoint is \(O(n^2)\). After discretizing, two Fenwick trees query “how many smaller on the left / larger on the right” in \(O(\log n)\): every value starts in the right tree and moves left as \(i\) advances.

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

Thinking

Fenwick trees need discretization. Memoizing by index, remaining picks, and direction counts increasing or decreasing completions: from \(i\) we try every later \(j\) that respects the order. \(O(n)\) states and \(O(n)\) transitions match the first method's bound, expressed as recursion.

 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