1395. 统计作战单位数

Difficulty: ** 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 中的元素都是唯一的 **

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：

1
2
3

输入：rating = [2,5,3,4,1]
输出：3
解释：我们可以组建三个作战单位 (2,3,4)、(5,4,1)、(5,3,1) 。

示例 2：

1
2
3

输入：rating = [2,1,3]
输出：0
解释：根据题目条件，我们无法组建作战单位。

示例 3：

1
2

输入：rating = [1,2,3,4]
输出：4

提示：

  • n == rating.length
  • 3 <= n <= 1000
  • 1 <= rating[i] <= 10^5
  • rating 中的元素都是唯一的

Solution

Language: golang

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

 
type  Tree []int

func lowbit(u int) int {
    return int(u & -u)
}
func (a Tree) Add(u ,c int) {
    for ;u<len(a);u+= lowbit(u) {
        a[u] += c
    }

}

func (a Tree) QueryCnt(u int) int {
    var ans int
    for ;u>0;u -= lowbit(u) {
        ans += a[u]
    }
    return ans
}

func set0(a[]int) {
    for i:=0;i<len(a);i++ {
        a[i] = 0
    }
}
func numTeams(rating []int) int {
    var n = len(rating)
    
    var leftSmall = make([]int,n+10)
    var leftBig = make([]int,n+10)
    
    var maxT = -1
    for _,v := range rating {
        //记录最大的 那个元素
        maxT = int( math.Max(float64(v),float64(maxT)) ) + 1
    }
    var tree = make(Tree,maxT+10)
    for i:=0;i<n;i++ {
        //统计左边比 rating[i] 小的那个数
        leftSmall[i] = tree.QueryCnt(rating[i]-1)
        //统计左边 大于 rating[i] 的个数
        leftBig[i] = tree.QueryCnt(maxT) - tree.QueryCnt(rating[i])
        //将当前元素记录到树中
        tree.Add(rating[i],1)
    }

    var ans int
    set0(tree)
    for i:=n-1;i>=0;i-- {
        var rightBig = tree.QueryCnt(maxT) - tree.QueryCnt(rating[i])
        var rightSmall = tree.QueryCnt(rating[i]-1)
        ans += rightBig * leftSmall[i]
        ans += rightSmall * leftBig[i]
        tree.Add(rating[i],1)
    }
    return ans
 

}