【三元组枚举中点】【树状数组】个人练习-Leetcode-1395. Count Number of Teams

题目链接：https://leetcode.cn/problems/count-number-of-teams/description/

题目大意：给一个数组rating[]，求符合以下任一条件的三元组i, j, k的个数

• rating[i] < rating[j] < rating[k]
• rating[i] > rating[j] > rating[k]

其实就是递增和递减。

思路：暴力枚举当然不太行。那么怎么搞定【三元组】呢？【三元组】中，中间的点总是特殊的，可以考虑枚举中点j，然后再在左边枚举i，右边枚举k，找出两边的大于和小于中间点的数量，然后【左小✖️右大 + 左大✖️右小】就是答案。这个枚举i和枚举k是同一层循环内的，因此时间复杂度是$O(N^2)$

完整代码

class Solution {
public:
    int numTeams(vector<int>& 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;
    } 
};

看了题解还有用树状数组的写法。树状数组建议看这个视频（https://www.bilibili.com/video/BV1ce411u7qP/）了解下，就能明白三个相关函数lowbit()，add()，query(()的作用。

但知道树状数组了，该怎么应用到这个题目呢？题目里可以作为的树状数组arr[]是什么呢？假设有个桶数组bk[]，为1表示这个下标的数字出现过，为0表示没出现过。然后对这个桶数组求前缀和得到一个数组，这个数组就是arr[]。在遍历rating[]的时候，每次遍历都会更新这个arr[]，这样就可以知道，在某个位置j左边小于rating[j]的数目，也就是潜在的iless。对于k，则倒过来再算一遍。由此可以记录iless, imore, kless, kmore。用和前面相同的方法计算即可。

完整代码

class Solution {
public:
    static constexpr int MAXN = 1001;
    int arr[MAXN];
    vector<int> disc;
    vector<int> iless, imore, kless, kmore;

    // aux func
    inline int lowbit(int x) {
        return x & (-x);
    }

    // add val to arr[pos]
    void add(int pos, int val) {
        while (pos < MAXN) {
            arr[pos] += val;
            pos += lowbit(pos);
        }
    }

    int query(int pos) {
        int res = 0;
        while (pos > 0) {
            res += arr[pos];
            pos -= lowbit(pos);
        }
        return res;
    }

    int numTeams(vector<int>& rating) {
        disc = rating;
        disc.emplace_back(-1);
        sort(disc.begin(), disc.end());
        auto getId = [&] (int target) {
            return lower_bound(disc.begin(), disc.end(), target) - disc.begin();
        };
        int n = rating.size();
        int ans = 0;
        iless.resize(n);
        imore.resize(n);
        kless.resize(n);
        kmore.resize(n);

        // forward
        for (int j = 0; j < n; j++) {
            auto id = getId(rating[j]);
            iless[j] = query(id);
            imore[j] = query(MAXN-1) - query(id);
            add(id, 1);
        }

        // reset arr to zero
        memset(arr, 0, sizeof arr);
        // backward
        for (int j = n-1; j >= 0; j--) {
            auto id = getId(rating[j]);
            kless[j] = query(id);
            kmore[j] = query(MAXN-1) - query(id);
            add(id, 1);
        }

        for (int i  = 0; i < n; i++) {
            ans += iless[i] * kmore[i] + imore[i] * kless[i];
        }
        
        return ans;
    } 
};