Question

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

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

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

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

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

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

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