3072. 将元素分配到两个数组中 II

思路：将所有非重复的元素进行排序离散化后，分别开2个树状数组记录2个组中各个元素出现的次数；接着遍历原数组，根据二分查找到当前元素在离散化数组中的下标，对比2个树状数组的返回值的大小来判断存入哪个数组中。

class BitTree {
private:
    const int n;
    vector<int> val;
public:
    BitTree(int _n) : n(_n), val(n + 1, 0) {}

    inline int lowbit(int i) {
        return i & -i;
    }

    void add(int i) {
        while(i <= n) {
            val[i]++;
            i += lowbit(i);
        }
    }
    
    int query(int i) {
        int ans = 0;
        while(i) {
            ans += val[i];
            i -= lowbit(i);
        }
        return ans;
    }
};


class Solution {
public:
    vector<int> resultArray(vector<int>& nums) {
        int n = nums.size();
        vector<int> sorted = nums;
        sort(sorted.begin(), sorted.end());
        sorted.erase(unique(sorted.begin(), sorted.end()), sorted.end());

        vector<int> a{nums[0]}, b{nums[1]};
        BitTree b1(n + 1), b2(n + 1);
        b1.add(ranges::lower_bound(sorted, nums[0]) - sorted.begin() + 1);
        b2.add(ranges::lower_bound(sorted, nums[1]) - sorted.begin() + 1);

        for(int i = 2;i < n;i++) {
            int index = ranges::lower_bound(sorted, nums[i]) - sorted.begin() + 1;
            int gc1 = a.size() - b1.query(index);
            int gc2 = b.size() - b2.query(index);
            if(gc1 > gc2 || gc1 == gc2 && a.size() <= b.size()) {
                a.push_back(nums[i]);
                b1.add(index);
            } else {
                b.push_back(nums[i]);
                b2.add(index);
            }
        }
        a.insert(a.end(), b.begin(), b.end());
        return a;
    }
};