题目描述

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

请你返回按上述条件可以组建的作战单位数量。每个士兵都可以是多个作战单位的一部分。

样例

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

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

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

限制

• n == rating.length
• 1 <= n <= 200
• 1 <= rating[i] <= 10^5

算法1

(暴力枚举) $O(n^3)$

1. 直接暴力枚举三个下标即可。

时间复杂度

• 三层循环，时间复杂度为 $O(n^3)$。

空间复杂度

• 仅需要常数的额外空间

C++ 代码

class Solution {
public:
    int numTeams(vector<int>& rating) {
        int n = rating.size();
        int ans = 0;

        for (int i = 0; i < n; i++)
            for (int j = i + 1; j < n; j++)
                for (int k = j + 1; k < n; k++) {
                    if (rating[i] < rating[j] && rating[j] < rating[k])
                        ans++;
                    else if (rating[i] > rating[j] && rating[j] > rating[k])
                        ans++;
                }

        return ans;
    }
};

算法2

(树状数组) $O(M + n^2 \log M)$

1. 枚举 i，然后从 i + 1 开始枚举 j。枚举 j 的过程中，用树状数组记录 [i + 1, j - 1] 之间的数字所在位置的前缀和。
2. 对于 rating[i] 和 rating[j]，如果前者小于后者，则答案累计 query(rating[j] - 1) - query(rating[i])，反之亦然。

时间复杂度

• 假设 $M$ 等于最大数字。
• 两层循环，树状数组的修改时间复杂度为 $O(\log M)$。
• 初始化树状数组需要 $O(M)$ 的时间。
• 总时间复杂度为 $O(M + n^2 \log M)$。
• 通过离散化，可以将时间复杂度优化到 $O(n^2 \log n)$。

空间复杂度

• 需要额外 $O(M)$ 的空间存储树状数组。

C++ 代码

class Solution {
public:
    const static int M = 100001;
    vector<int> f;

    int query(int x) {
        int tot = 0;
        for (; x; x -= x & -x)
            tot += f[x];

        return tot;
    }

    void add(int x, int y) {
        for (; x < M; x += x & -x)
            f[x] += y;
    }

    int numTeams(vector<int>& rating) {
        int n = rating.size();
        int ans = 0;
        f.resize(M, 0);

        for (int i = 0; i < n; i++) {
            for (int j = i + 1; j < n; j++) {
                if (rating[j] > rating[i]) {
                    ans += query(rating[j] - 1) - query(rating[i]);
                } else if (rating[j] < rating[i]) {
                    ans += query(rating[i] - 1) - query(rating[j]);
                }
                add(rating[j], 1);
            }
            for (int j = i + 1; j < n; j++)
                add(rating[j], -1);
        }

        return ans;
    }
};

算法3

(枚举) $O(n^2)$

1. 枚举中间的位置 $i$，然后找 $i$ 左边小于 rating[i] 和大于 rating[i] 的个数 l1 和 l2，右边同理求出 r1 和 r2。
2. 则 l1 * r2 + l2 * r1 就是中间位置 $i$ 所做的贡献。

时间复杂度

• 两层循环，时间复杂度为 $O(n^2)$。
• 此算法仍可以通过树状数组优化到 $O(M + n \log M)$，通过离散化可以进一步优化到 $O(n \log n)$。

空间复杂度

• 仅需要常数的额外空间。

C++ 代码

class Solution {
public:
    int numTeams(vector<int>& rating) {
        int n = rating.size();
        int ans = 0;

        for (int i = 0; i < n; i++) {
            int l1 = 0, r1 = 0, l2 = 0, r2 = 0;
            for (int j = i - 1; j >= 0; j--)
                if (rating[j] < rating[i])
                    l1++;
                else if (rating[j] > rating[i])
                    l2++;

            for (int j = i + 1; j < n; j++)
                if (rating[j] < rating[i])
                    r1++;
                else if (rating[j] > rating[i])
                    r2++;

            ans += l1 * r2 + l2 * r1;
        }

        return ans;
    }
};

算法4

(树状数组优化算法 3) $O(n \log n)$

1. 在算法 3 的基础上，用树状数组求 l1, l2, r1, r2。
2. 同时配合上离散化，将数域缩小到 $O(n)$。

时间复杂度

• 初始化需要 $O(n \log n)$ 的时间离散化，需要 $O(n \log n)$ 的时间提前维护一个树状数组。
• 树状数组每次查询和更新仅需要 $O(\log n)$ 的时间。
• 故总时间复杂度为 $O(n \log n)$。

空间复杂度

• 需要额外 $O(n)$ 的空间离散化和存储树状数组。

C++ 代码

class Solution {
public:
    int m;

    int query(const vector<int> &f, int x) {
        int tot = 0;
        for (; x; x -= x & -x)
            tot += f[x];

        return tot;
    }

    void add(vector<int> &f, int x, int y) {
        for (; x <= m; x += x & -x)
            f[x] += y;
    }

    int numTeams(vector<int>& rating) {
        int n = rating.size();
        int ans = 0;

        vector<int> s(rating);
        sort(s.begin(), s.end());
        m = unique(s.begin(), s.end()) - s.begin();
        s.resize(m);
        for (int i = 0; i < n; i++)
            rating[i] = lower_bound(s.begin(), s.end(), rating[i]) - s.begin() + 1;

        vector<int> f(m + 1, 0), g(m + 1, 0);

        for (int i = 0; i < n; i++)
            add(g, rating[i], 1);


        for (int i = 0; i < n; i++) {
            int l1 = query(f, rating[i] - 1);
            int l2 = i - query(f, rating[i]);
            int r1 = query(g, rating[i] - 1);
            int r2 = n - i - query(g, rating[i]);

            ans += l1 * r2 + l2 * r1;

            add(f, rating[i], 1);
            add(g, rating[i], -1);
        }

        return ans;
    }
};