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