Welcome to Subscribe On Youtube

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$.

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$.

  • 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;
        }
    }
    
  • 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;
        }
    };
    
  • 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
    
    
  • 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
    }
    
  • 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;
    }
    
    
  • 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;
        }
    }
    
  • 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;
        }
    };
    
  • 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
    
    
  • 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
    }
    
  • 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;
    }
    
    
  • 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';
    
    
  • /**
     * @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;
    };
    
    

All Problems

All Solutions