LeetCode 1395 统计作战单位数——树状数组初学

题目：

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
请你返回按上述条件可以组建的作战单位数量。每个士兵都可以是多个作战单位的一部分。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/count-number-of-teams

思路

根据官方给的题解，有三个方法

方法一：暴力破解

套用三重循坏进行枚举，显然没有什么价值，考虑剩下两个方法。

方法二：枚举中间点

枚举(i, j, k)中的j，并分别统计数组中在j两侧中大于或小于j的数据个数lmore,lless,rmore,rless,具体求的过程也是很简单，就不解释了

易得答案为lmore * rless + lless * rmore

class Solution:
    def numTeams(self, rating: List[int]) -> int:
        res = 0
        for j in range(len(rating)):
            lmore = lless = 0
            rmore = rless = 0
            for i in range(j):
                if rating[i] > rating[j]:
                    lmore += 1
                else:
                    lless += 1
            for k in range(j + 1, len(rating)):
                if rating[k] > rating[j]:
                    rmore += 1
                else:
                    rless += 1
            res +=  lmore * rless + lless * rmore
        return res

方法三：离散化树状数组

有几个前提概念需要讲，一个是树状数组和lowbit，这里参考：树状数组 & lowbit() - liubilan - 博客园

还有就是桶和离散化思想(具体我也解释不清，反正这么叫就对了),具体可以参考：力扣

具体就是优化了方法二中计算lmore,lless,rmore,rless的时间效率

具体代码：

(思路写的并不是很好，希望以后的自己还能看懂)

class Solution:
    def numTeams(self, rating: List[int]) -> int:
        import bisect
        def getId(n):#获取n在rating按大小排序后的逻辑位置
            return bisect.bisect_left(l, n) + 1

        def lowbit(x):#计算lowbit，用于维护树状数组
            return x&-x

        def get(p):#取得逻辑位置p的数的前缀和
            #通过树状数组查询的方式获取以优化时间效率
            r = 0
            while p > 0:
                r += c[p]
                p -=  lowbit(p)
            return r

        def add(p):#更新逻辑位置p的数和受p影响的前缀和的状态
            #将桶中位置p的数字更新为1，代表该数字出现，即代表在rating中下一个数的左边有逻辑位置为p的数(l[p]),同时更新受位置p影响的前缀和
            while p <= length:
                c[p] += 1
                p += lowbit(p)


        length = len(rating)
        c = [0 for i in range(length + 1)]  #桶数组，记录前缀和
        l = rating[:]   #排序后的rating数组，用于获取逻辑位置
        l.sort()
        Lless = []
        Lmore = []
        Rless = []
        Rmore = []

        for i in range(length):#正向遍历求Lless和Lmore
            id = getId(rating[i])
            num = get(id)
            Lless.append(num)
            Lmore.append(i - num)
            add(id)

        c = [0 for i in range(length + 1)] 
        for i in range(length - 1, -1, -1):#反向遍历求Rless和Rmore
            id = getId(rating[i])
            num = get(id)
            Rless.append(num)
            Rmore.append(length - 1 - i - num)
            add(id)
        Rless.reverse()
        Rmore.reverse()

        res = 0
        for i in range(length):
            res += Lmore[i] * Rless[i] + Lless[i] * Rmore[i]

        return res