1395. Count Number of Teams
给你n个士兵，每个人都有唯一的评价
需要将3个士兵编成一个队伍
须满足下列条件
1.3个士兵的index分别是i、j、k，i<j<k
2.rating[i]<rating[j]<rating[k] or rating[i]>rating[j]>rating[k]
1个士兵可以编在多个队伍里
请问总共会有几个队伍
思路:
我一开始用n^3，后来想想不对用n^2
接着就去看答案用binary index tree
就建立一个矩阵arr，其中arr[i]是第i个士兵的rating值是第几小的
接着再从0开始到n-1边计算目前这个士兵
左边有几个比他小的、右边有几个比他大的
计算完后把这个士兵丢到binary index tree
接着再做一次，这次要计算目前这个士兵
左边有几个比他大的、右边有几个比他小的
最后就可以得到答案了
golang code :
type soldier struct {
idx int
val int
}
func numTeams(rating []int) int {
n := len(rating)
soldiers := make([]soldier, n)
for key, val := range rating {
soldiers[key].idx = key
soldiers[key].val = val
}
slices.SortFunc(soldiers, func(i, j soldier) int {
return i.val - j.val
})
for key, val := range soldiers {
rating[val.idx] = key + 1
}
calculate := func(ascending bool) int {
bits := make([]int, n+1)
teams := 0
for i := 0; i < n; i++ {
left, right, rank := 0, 0, rating[i]
if ascending {
left = getsum(bits, rank-1)
right = n - rank - (getsum(bits, n) - getsum(bits, rank))
} else {
left = getsum(bits, n) - getsum(bits, rank)
right = rank - 1 - getsum(bits, rank-1)
}
teams += left * right
update(bits, rank, 1)
}
return teams
}
return calculate(true) + calculate(false)
}
func getsum(bits []int, rank int) int {
sum := 0
for rank > 0 {
sum += bits[rank]
rank -= (rank & -rank)
}
return sum
}
func update(bits []int, rank, val int) {
n := len(bits)
for rank < n {
bits[rank] += val
rank += (rank & -rank)
}
}