1395. 统计作战单位数

题目：
leetcode.cn/problems/co…

算法：
方法一：枚举
iless：i<j,rating[i]<rating[j]的i的个数
igreat：i<j,rating[i]>rating[j]的i的个数
kless,k>j,rating[k] < rating[j]的k的个数
kmore,k>j,rating[k] > rating[j]的k的个数
对每一个j,符合条件的作战单位个数为：iless * kgreat + igreat * iless

func numTeams(rating []int) int {
    n := len(rating)
    ans := 0
    for j := 1; j < n - 1; j ++ {
        iless, igreat, kless, kgreat := 0, 0, 0, 0
        for i := 0; i < j; i ++ {
            if rating[i] < rating[j] {
                iless ++
            }
            if rating[i] > rating[j] {
                igreat ++
            }
        }
        for k := j + 1; k < n; k ++ {
            if rating[j] < rating[k] {
                kgreat ++
            }
            if rating[j] > rating[k] {
                kless ++
            }
        }
        ans = ans + iless * kgreat + igreat * kless
    }
    return ans
}

方法二：树状数组
我们需要对每一个index，统计出iless,igreat,kless,kgreat。结果自然就能求出来了。总体思路：

iless,igeat是i左侧，小于i和大于i的数的个数。我们可以用BIT[i]维护小于i的数的个数：

1. 从左往右遍历rating，先统计小于BIT[i]的数的个数，iless就是BIT[i]，igreat是n - iless。然后对BIT[i]加一（对BIT[i~n]进行更新）。
2. 重建BIT，从右往左遍历rating，遍历到位置i，kless = BIT[n]，kgreat = n - iless

func numTeams(rating []int) int {
    n := len(rating)
    maxN := n + 2
    iless, igreat := make([]int, n), make([]int, n)
    kless, kgreat := make([]int, n), make([]int, n)
    ans := 0
    
    tmp := make([]int, n + 1)
    tmp[n] = -1
    copy(tmp, rating)
    sort.Ints(tmp)
    // 离散化，因为rating[i]的值范围太大，映射到len(rating)的范围
    // 返回rating中小于val的元素个数
    // id要从1开始[1,n]，因为0 & 数都为0，导致lowbit一直返回0，add陷入死循环
    var getID func(val int) int
    getID = func(val int) int{
        left, right := 0, len(tmp) - 1
        for left < right {
            mid := (left + right) / 2
            if tmp[mid] < val {
                left = mid + 1
            } else {
                right = mid
            }
           
        }
        return left
    }
    // BIT[i]表示i左侧，小于id的元素个数
    BIT = make([]int, maxN)
    for i := 0; i < n; i ++ {
        id := getID(rating[i])
        // rating[i]是唯一的，所以id是唯一的，又因为最后在进行add，
        // 所以iless[i]表示小于的元素个数
        iless[i] = get(id)
        igreat[i] = get(n + 1) - get(id)
        add(id, 1)
    }

    BIT = make([]int, maxN)
    for i := n - 1; i >= 0 ; i -- {
        id := getID(rating[i])
        kless[i] = get(id)
        kgreat[i] = get(n + 1) - get(id)
        add(id, 1)
    }

    for i := 0; i < n; i ++ {
        ans = ans + iless[i] * kgreat[i] + igreat[i] * kless[i]
    }
    return ans
}

var BIT []int

func add(index, val int) {
    for index < len(BIT) {
        BIT[index] = BIT[index] + val
        index = index + lowBit(index)
    }
}

func get(index int) int {
    ans := 0 
    for index > 0 {
        ans = ans + BIT[index]
        index = index - lowBit(index)
    }
    return ans
}

// 返回将x最低位的1左侧截断，得到的数
func lowBit(x int) int {
    return x & (-x)
}