LeetCode: 1395. 统计作战单位数

目录

1. 解法一：枚举中点

2. 解法二：树状数组 + 离散化 优化解法一

---

原题链接：1395. 统计作战单位数 - 力扣（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 中的元素都是唯一的

1. 解法一：枚举中点

要找到有多少个满足条件的三元组，我们可以枚举三元组的中间位置，通过统计中间位置两侧满足条件的元素的数量来解决这个问题。

对于满足条件的三元组有以下两种可能：

1：随着下标的增长，三元组严格递增。在枚举中间位置时，我们就需要统计出中间位置左侧比中间位置元素小的元素个数，再统计出中间位置右侧比中间位置大的元素个数。两者的乘积就是这个中间位置，三元组严格递增的数目。

例如：2,1,3,5,4  当我们枚举到 3 时，统计 3 的左侧小于 3 的元素个数：2，  统计 3 的右侧大于 3 的元素个数：2。那么，以 3 为中间位置，严格递增的三元组数目就是 2 * 2 = 4 (一个简单的组合问题)。

2：随着下标的增长，三元组严格递减。原理和上面是一样的哇！我们需要统计出中间位置左侧比中间位置元素大的元素个数，再统计出中间位置右侧比中间位置小的元素个数。两者的乘积就是这个中间位置，三元组严格递减的数目。

例如：6,5,3,4,2  当我们枚举到 3 时，统计 3 的左侧大于 3 的元素个数：2，  统计 3 的右侧小于 3 的元素个数：1。那么，以 3 为中间位置，严格递减的三元组数目就是 2 * 1 = 2 (一个简单的组合问题)。

重复上述操作，枚举数组中的所有可能的中间位置，计算每个位置的满足条件的三元组个数，求和即可！

 

class Solution {
public:
    int numTeams(vector& rating) {
        int n = rating.size();
        int ans = 0;
        for (int j = 1; j < n - 1; ++j) {
            //左侧小于      左侧大于
            int iless = 0, imore = 0;
            //右侧小于      右侧大于
            int kless = 0, kmore = 0;
            for (int i = 0; i < j; ++i) {
                if (rating[i] < rating[j]) {
                    ++iless;
                }
                else if (rating[i] > rating[j]) {
                    ++imore;
                }
            }
            for (int k = j + 1; k < n; ++k) {
                if (rating[k] < rating[j]) {
                    ++kless;
                }
                else if (rating[k] > rating[j]) {
                    ++kmore;
                }
            }
            ans += iless * kmore + imore * kless;
        }
        return ans;
    }
};

时间复杂度分析：我们在枚举中间位置的时候需要向左和向右遍历数组找到较大和较小元素的个数。因此时间复杂度是：O(N * N)。空间复杂度: O(1)。

2. 解法二：树状数组 + 离散化 优化解法一

在解法二中，我们枚举每个中间位置时，都需要向左和向右找比中间位置大和小的元素个数。这个过程能否优化呢？

我们可以创建一个数组arr, 初始化为0，遍历到一个数字num，就让arr[num] = 1。

如上图：我们遍历到 3 时 2，5，已经出现过了，那么arr[2] 和 arr[5] 就是 1， 

我们要找左侧比 3 小的元素，只需要知道下标为 3 的位置的前缀和即可(如图可知前缀和为1)

我们要找左侧比 3 大的元素，只需要知道整个数组的前缀和(如图可知前缀和为2)，然后减去下标为 3 的位置的前缀和(如图可知前缀和为 1 ) 即可：2 - 1 = 1。

计算出结果之后，再让 arr[3] = 1, 代表 3 出现过了！

通过这样的方法我们就可以在较快时间内求出，左侧较小的元素个数，左侧较大的元素个数。

求左侧较小元素的个数与左侧较大元素的个数需要从左向右遍历原数组，求右侧较小元素的个数与右侧较大元素的个数就需要从右向左遍历原数组啦！ 

这里的这个 arr 数组如果用普通的前缀和数组，时间复杂度还是没有变，因为我们需要添加新的数到前缀和数组。涉及到单点更新的前缀和维护自然就是树状数组啦！

我们还注意到，原数组的大小只有 1000，但是数组中的元素大小有 10^5 级别的，因此我们可以考虑使用整数的离散化，减小树状数组的高度，优化单点修改和区间查询的效率。

整数保序的离散化(C/C++)_c++离散化-CSDN博客

这是离散化的讲解，树状数组与线段树还没写！如果不熟悉可以用其他大佬的博客复习复习。

 

const int N = 1010; //离散化之后的最大数据范围
int tr[N]; //树状数组
class Solution {
public:
    //树状数组的三个函数
    int lowbit(int x)
    {
        return x & -x;
    }
    //这道题目的add仅仅是加一，就不用多谢一个参数啦
    void add(int x)
    {
        for(int i = x; i <= N; i += lowbit(i))
            tr[i] += 1;
    }

    int query(int x)
    {
        int res = 0;
        for(int i = x; i > 0; i -= lowbit(i))
            res += tr[i];
        return res;
    }
    //离散化需要用到的二分查找，有库函数我不用，唉就是玩儿
    int BinarySearch(vector& nums, int target)
    {
        int l = 0, r = nums.size() - 1;
        while(l <= r)
        {
            int mid = (l + r) >> 1;
            if(nums[mid] < target) l = mid + 1;
            else if(nums[mid] > target) r = mid - 1;
            else return mid + 1;
        }
        return -1;
    }

    int numTeams(vector& rating) {
        int n = rating.size();
        //copy数组用来离散化
        vector copy(rating);
        //题目说了没有重复元素，不需要去重
        sort(copy.begin(), copy.end());
        //记录左右两侧较大和较小的元素个数
        vector iless(n);
        vector ibig(n);
        vector kless(n);
        vector kbig(n);
        
        int index = 0;
        //初始化树状数组
        memset(tr, 0, sizeof tr);
        //从前向后遍历找左侧
        for(int i = 0; i < n; i++)
        {
            //查找离散化之后的结果
            index = BinarySearch(copy, rating[i]);
            //左侧较小
            iless[i] = query(index);
            //右侧较大
            ibig[i] = query(1000) - query(index);
            //加入新的数
            add(index);
        }
        //初始化
        memset(tr, 0, sizeof tr);
        //从后往前遍历找右侧
        for(int i = n - 1; i >= 0; i--)
        {
            //同上
            index = BinarySearch(copy, rating[i]);
            kless[i] = query(index);
            kbig[i] = query(1000) - query(index);
            add(index);
        }
        //计算结果！
        int ans = 0;
        for(int i = 0; i < n; i++)
        {
            ans += iless[i] * kbig[i] + ibig[i] * kless[i];
        }
        return ans;
    }
};

 时间复杂度分析：树状数组单点修改与区间查询的时间复杂度都是 O(logN), 因此最终的时间复杂度：O(N*logN)，空间复杂度：O(N).