夯实算法-统计作战单位数

题目：LeetCode

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

解题思路

遍历原数组，对于rating[i]，使用二分查找确定离散化后的位置id。 遍历过程中，不断将元素入桶，也就是在BIT.arr中，id对应的位置设为1，表示此元素已经存在。

preLess[i]，代表在rating数组中，在rating[i]前面，且比rating[i]小的元素个数； preMore[i]，代表在rating数组中，在rating[i]前面，且比rating[i]大的元素个数；

需要明白的一点是，此时是顺序遍历，遍历到i时，0到i-1位置的元素全部已经入桶， 所以能够计算preLess[i]和preMore[i]；

又因为我们需要不断将元素入桶（也就是单点更新），同时查询位置i前面或后面元素的个数（也就是区间查询）， 自然使用树状数组这一数据结构。 树状数组做单点更新，区间查询的时间复杂度均是O(logn)；

代码实现

int n;
public int numTeams(int[] rating) {
    n = rating.length;
    int disc[] = discretization(rating);

    BIT bit1 = new BIT(n);

    int preLess[] = new int[n];
    int preMore[] = new int[n];
    for (int i = 0; i < n; i++) {
        int id = get_Id(disc, rating[i]);
        preLess[i] = bit1.query(id); //查询前面有多少个元素
        preMore[i] = bit1.query(n) - bit1.query(id); //查询后面有多少个元素
        bit1.update(id, 1);
    }

    BIT bit2 = new BIT(n);

    int afterLess[] = new int[n];
    int afterMore[] = new int[n];
    for (int i = n - 1; i >= 0; i--) {
        int id = get_Id(disc, rating[i]);
        afterLess[i] = bit2.query(id);
        afterMore[i] = bit2.query(n) - bit2.query(id);
        bit2.update(id, 1);
    }

    int cnt = 0;
    for (int i = 1; i < n - 1; i++)
        cnt += preLess[i] * afterMore[i] + preMore[i] * afterLess[i];

    return cnt;
}

//对原数组做离散化处理，提高执行效率，又因为题目明确说明没有重复元素，无须去重
public int[] discretization(int[] rating) {
    int disc[] = new int[n];
    for (int i = 0; i < n; i++) disc[i] = rating[i];
    Arrays.sort(disc);
    return disc;
}

public int get_Id(int[] disc, int x) {
    int left = 0, right = n - 1;
    while (left < right) {
        int mid = (left + right) / 2;
        if (disc[mid] < x)
            left = mid + 1;
        else
            right = mid;
    }
    return left + 1; //id 是树状数组的下标，故加一；
}
}


class BIT {
    int n;
    int arr[];

    public BIT(int len) {
        n = len;
        arr = new int[n + 1]; //arr数组的大小要比原数组大1
    }

    public void update(int pos, int x) {
        while (pos <= n) {
            arr[pos] += x;
            pos += low_bit(pos);
        }
    }

    public int low_bit(int x) {
        return x & (-x);
    }

    public int query(int pos) {
        int ans = 0;
        while (pos > 0) {
            ans += arr[pos];
            pos -= low_bit(pos);
        }
        return ans;
    }

运行结果

复杂度分析

• 空间复杂度：$O(1)$
• 时间复杂度：$O(logn)$

在掘金(JUEJIN)  一起分享知识, Keep Learning!