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;
}
};