1395. 统计作战单位数

Difficulty
Medium
Tags
树状数组
URL
https://leetcode.cn/problems/count-number-of-teams/
Star
n 名士兵站成一排。每个士兵都有一个 独一无二 的评分 rating 。
每 3 个士兵可以组成一个作战单位，分组规则如下：
  • 从队伍中选出下标分别为 i、j、k 的 3 名士兵，他们的评分分别为 rating[i]、rating[j]、rating[k]
  • 作战单位需满足： rating[i] < rating[j] < rating[k] 或者 rating[i] > rating[j] > rating[k] ，其中  0 <= i < j < k < n
请你返回按上述条件可以组建的作战单位数量。每个士兵都可以是多个作战单位的一部分。
示例 1：
输入：rating = [2,5,3,4,1]
输出：3
解释：我们可以组建三个作战单位 (2,3,4)、(5,4,1)、(5,3,1) 。
示例 2：
输入：rating = [2,1,3]
输出：0
解释：根据题目条件，我们无法组建作战单位。
示例 3：
输入：rating = [1,2,3,4]
输出：4
提示：
  • n == rating.length
  • 3 <= n <= 1000
  • 1 <= rating[i] <= 10^5
  • rating 中的元素都是唯一的

法1 树状数组

思路
以 nums[i] 为中心的 满足条件的个数等于 num[i] 左边小于 num[i] 的个数 * num[i] 右边大于 num[i] 的个数 + num[i] 左边大于 num[i] 的个数 * num[i] 右边大小于 num[i] 的个数。
所以最终答案就是：遍历 nums, 将所有以nums[i]为中心的满足条件的个数相加即可。
题解
class NumArray:
    def __init__(self, n):
        self.n = n
        self.tree = [0] * n

    def lowbit(self, x):
        return x & (-x)

    def add(self, idx, val=1):
        idx += 1
        while idx < len(self.tree):
            self.tree[idx] += val
            idx += self.lowbit(idx)

    def prefix(self, idx):
        idx += 1
        sum_ = 0
        while idx > 0:
            sum_ += self.tree[idx]
            idx -= self.lowbit(idx)
        return sum_
    
    def rangesum(self, left, right=-1):
        if right == -1:
            right = len(self.tree)-2
        return self.prefix(right) - self.prefix(left-1)


class Solution:
    def numTeams(self, rating: List[int]) -> int:
        if len(rating) < 3:
            return 0

        # 压缩
        unique = sorted(list(set(rating)))
        hash_map = {v:i for i, v in enumerate(unique)}

        left_small = [0 for _ in rating]
        left_big = [0 for _ in rating]
        right_small = [0 for _ in rating]
        right_big = [0 for _ in rating]
        
        # 统计 左边比自己小的 和 左边比自己大的
        T1 = NumArray(unique[-1]+2)
        for idx in range(len(rating)):
            num = hash_map[rating[idx]]
            left_small[idx] = T1.prefix(num-1)
            left_big[idx] = T1.rangesum(left=num+1)
            T1.add(num)

        # 统计 右边比自己小的 和 右边边比自己大的
        T2 = NumArray(unique[-1]+2)
        for idx in range(len(rating)-1, -1, -1):
            num = hash_map[rating[idx]]
            right_small[idx] = T2.prefix(num-1)
            right_big[idx] = T2.rangesum(left=num+1)
            T2.add(num)
        
        # 得到答案
        res = 0
        for i in range(len(rating)):
            res += left_big[i] * right_small[i] + left_small[i] * right_big[i]

        return res